<commit_before>/* SPDX-License-Identifier: BSD-2-Clause */

#ifndef _KMALLOC_C_

   #error This is NOT a header file and it is not meant to be included

#endif

/*
 * NOTE: the trick to make the small heap to work well without the number of
 * small heaps to explode is to allow it to allocate just a small fraction of
 * its actual size, like 1/16th.
 */
#define SMALL_HEAP_MAX_ALLOC (SMALL_HEAP_SIZE / 16 - 1)

typedef struct {

   list_node node;
   list_node avail_node;
   kmalloc_heap heap;

} small_heap_node;

static kmalloc_small_heaps_stats shs;
static list small_heaps_list;
static list avail_small_heaps_list;

static inline void
add_in_avail_list(small_heap_node *node)
{
   ASSERT(node->heap.mem_allocated < node->heap.size);
   shs.not_full_count++;

   if (shs.not_full_count > shs.peak_not_full_count)
      shs.peak_not_full_count = shs.not_full_count;

   list_add_tail(&avail_small_heaps_list, &node->avail_node);
}

static inline void
remove_from_avail_list(small_heap_node *node)
{
   ASSERT(!list_node_is_empty(&node->avail_node));

   ASSERT(
      node->heap.mem_allocated == node->heap.size ||
      node->heap.mem_allocated == SMALL_HEAP_MD_SIZE
   );

   ASSERT(shs.not_full_count > 0);

   list_remove(&node->avail_node);
   shs.not_full_count--;

   DEBUG_ONLY(list_node_init(&node->avail_node));
}

static inline void
register_small_heap_node(small_heap_node *node)
{
   ASSERT(node->heap.mem_allocated < node->heap.size);

   list_add_tail(&small_heaps_list, &node->node);
   shs.tot_count++;

   if (shs.tot_count > shs.peak_count)
      shs.peak_count = shs.tot_count;

   ASSERT(list_node_is_empty(&node->avail_node));
   add_in_avail_list(node);
}

static inline void
unregister_small_heap_node(small_heap_node *node)
{
   ASSERT(!list_node_is_empty(&node->avail_node));
   remove_from_avail_list(node);

   ASSERT(shs.tot_count > 0);

   list_remove(&node->node);
   shs.tot_count--;

   ASSERT(node->heap.mem_allocated == SMALL_HEAP_MD_SIZE);
   ASSERT(list_node_is_empty(&node->avail_node));
}

static small_heap_node *alloc_new_small_heap(void)
{
   small_heap_node *new_node;
   void *heap_data, *metadata;
   size_t actual_metadata_size;
   size_t small_heap_sz = SMALL_HEAP_SIZE;

   /*
    * NOTE: here KMALLOC_FL_DONT_ACCOUNT is used in other to avoid the "heavy"
    * kmalloc stats (when enabled) to account this allocation, as it's not
    * really a proper allocation: it's the creation of a small heap. Instead,
    * the allocation for its metadata is explicitly accounted (see below) since
    * some memory was actually consumed.
    */

   if (!(heap_data = general_kmalloc(&small_heap_sz, KMALLOC_FL_DONT_ACCOUNT)))
      return NULL;

   metadata = heap_data;
   new_node = kzmalloc(MAX(sizeof(small_heap_node), SMALL_HEAP_MAX_ALLOC + 1));

   if (!new_node) {
      kfree2(heap_data, SMALL_HEAP_SIZE);
      return NULL;
   }

   list_node_init(&new_node->node);
   list_node_init(&new_node->avail_node);

   bool success =
      kmalloc_create_heap(&new_node->heap,
                          (uptr)heap_data,
                          SMALL_HEAP_SIZE,
                          SMALL_HEAP_MBS,
                          0,
                          true,
                          metadata,
                          NULL,
                          NULL);

   if (!success) {
      kfree2(new_node, MAX(sizeof(small_heap_node), SMALL_HEAP_MAX_ALLOC + 1));
      kfree2(heap_data, SMALL_HEAP_SIZE);
      return NULL;
   }

   // Allocate heap's metadata inside the heap itself
   actual_metadata_size = new_node->heap.metadata_size;

   DEBUG_ONLY_UNSAFE(void *actual_metadata =)
      per_heap_kmalloc(&new_node->heap, &actual_metadata_size, 0);
   ASSERT(actual_metadata == metadata);

   if (KMALLOC_HEAVY_STATS)
      kmalloc_account_alloc(actual_metadata_size);

   register_small_heap_node(new_node);
   return new_node;
}

static void destroy_small_heap(small_heap_node *node)
{
   unregister_small_heap_node(node);
   kfree2((void *)node->heap.vaddr, SMALL_HEAP_SIZE);
   kfree2(node, MAX(sizeof(small_heap_node), SMALL_HEAP_MAX_ALLOC + 1));
}

static void *small_heaps_kmalloc(size_t *size, u32 flags)
{
   small_heap_node *new_node;
   small_heap_node *pos;
   void *ret;

   ASSERT(!is_preemption_enabled());
   ASSERT(*size <= (SMALL_HEAP_SIZE - SMALL_HEAP_MBS));

   list_for_each_ro(pos, &avail_small_heaps_list, avail_node) {

      if (pos->heap.size - pos->heap.mem_allocated < *size)
         continue;

      /* We've found a heap with (potentially) enough space */
      ret = per_heap_kmalloc(&pos->heap, size, flags);

      if (!ret)
         continue;

      /*
       * The alloc succeeded. If now the heap is full, remove it from the
       * 'avail' list.
       */

      if (pos->heap.mem_allocated == pos->heap.size)
         remove_from_avail_list(pos);

      return ret;
   }

   if (!(new_node = alloc_new_small_heap()))
      return NULL;<commit_msg>[kmalloc] Small heaps: more refactoring<commit_after>/* SPDX-License-Identifier: BSD-2-Clause */

#ifndef _KMALLOC_C_

   #error This is NOT a header file and it is not meant to be included

#endif

/*
 * NOTE: the trick to make the small heap to work well without the number of
 * small heaps to explode is to allow it to allocate just a small fraction of
 * its actual size, like 1/16th.
 */
#define SMALL_HEAP_MAX_ALLOC (SMALL_HEAP_SIZE / 16 - 1)

#define SMALL_HEAP_NODE_ALLOC_SZ \
   MAX(sizeof(small_heap_node), SMALL_HEAP_MAX_ALLOC + 1)

typedef struct {

   list_node node;
   list_node avail_node;
   kmalloc_heap heap;

} small_heap_node;

static kmalloc_small_heaps_stats shs;
static list small_heaps_list;
static list avail_small_heaps_list;

static inline small_heap_node *alloc_small_heap_node(void)
{
   return kzmalloc(SMALL_HEAP_NODE_ALLOC_SZ);
}

static inline void free_small_heap_node(small_heap_node *obj)
{
   kfree2(obj, SMALL_HEAP_NODE_ALLOC_SZ);
}

static inline void
add_in_avail_list(small_heap_node *node)
{
   ASSERT(node->heap.mem_allocated < node->heap.size);
   shs.not_full_count++;

   if (shs.not_full_count > shs.peak_not_full_count)
      shs.peak_not_full_count = shs.not_full_count;

   list_add_tail(&avail_small_heaps_list, &node->avail_node);
}

static inline void
remove_from_avail_list(small_heap_node *node)
{
   ASSERT(!list_node_is_empty(&node->avail_node));

   ASSERT(
      node->heap.mem_allocated == node->heap.size ||
      node->heap.mem_allocated == SMALL_HEAP_MD_SIZE
   );

   ASSERT(shs.not_full_count > 0);

   list_remove(&node->avail_node);
   shs.not_full_count--;

   DEBUG_ONLY(list_node_init(&node->avail_node));
}

static inline void
register_small_heap_node(small_heap_node *node)
{
   ASSERT(node->heap.mem_allocated < node->heap.size);

   list_add_tail(&small_heaps_list, &node->node);
   shs.tot_count++;

   if (shs.tot_count > shs.peak_count)
      shs.peak_count = shs.tot_count;

   ASSERT(list_node_is_empty(&node->avail_node));
   add_in_avail_list(node);
}

static inline void
unregister_small_heap_node(small_heap_node *node)
{
   ASSERT(!list_node_is_empty(&node->avail_node));
   remove_from_avail_list(node);

   ASSERT(shs.tot_count > 0);

   list_remove(&node->node);
   shs.tot_count--;

   ASSERT(node->heap.mem_allocated == SMALL_HEAP_MD_SIZE);
   ASSERT(list_node_is_empty(&node->avail_node));
}

static small_heap_node *alloc_new_small_heap(void)
{
   small_heap_node *new_node;
   void *heap_data, *metadata;
   size_t actual_metadata_size;
   size_t small_heap_sz = SMALL_HEAP_SIZE;

   /*
    * NOTE: here KMALLOC_FL_DONT_ACCOUNT is used in other to avoid the "heavy"
    * kmalloc stats (when enabled) to account this allocation, as it's not
    * really a proper allocation: it's the creation of a small heap. Instead,
    * the allocation for its metadata is explicitly accounted (see below) since
    * some memory was actually consumed.
    */

   if (!(heap_data = general_kmalloc(&small_heap_sz, KMALLOC_FL_DONT_ACCOUNT)))
      return NULL;

   ASSERT(small_heap_sz == SMALL_HEAP_SIZE);
   metadata = heap_data;
   new_node = alloc_small_heap_node();

   if (!new_node) {
      kfree2(heap_data, SMALL_HEAP_SIZE);
      return NULL;
   }

   list_node_init(&new_node->node);
   list_node_init(&new_node->avail_node);

   bool success =
      kmalloc_create_heap(&new_node->heap,
                          (uptr)heap_data,
                          SMALL_HEAP_SIZE,
                          SMALL_HEAP_MBS,
                          0,
                          true,
                          metadata,
                          NULL,
                          NULL);

   if (!success) {
      free_small_heap_node(new_node);
      kfree2(heap_data, SMALL_HEAP_SIZE);
      return NULL;
   }

   // Allocate heap's metadata inside the heap itself
   actual_metadata_size = new_node->heap.metadata_size;

   DEBUG_ONLY_UNSAFE(void *actual_metadata =)
      per_heap_kmalloc(&new_node->heap, &actual_metadata_size, 0);
   ASSERT(actual_metadata == metadata);

   if (KMALLOC_HEAVY_STATS)
      kmalloc_account_alloc(actual_metadata_size);

   register_small_heap_node(new_node);
   return new_node;
}

static void destroy_small_heap(small_heap_node *node)
{
   unregister_small_heap_node(node);
   kfree2((void *)node->heap.vaddr, SMALL_HEAP_SIZE);
   free_small_heap_node(node);
}

static void *small_heaps_kmalloc(size_t *size, u32 flags)
{
   small_heap_node *new_node;
   small_heap_node *pos;
   void *ret;

   ASSERT(!is_preemption_enabled());
   ASSERT(*size <= (SMALL_HEAP_SIZE - SMALL_HEAP_MBS));

   list_for_each_ro(pos, &avail_small_heaps_list, avail_node) {

      if (pos->heap.size - pos->heap.mem_allocated < *size)
         continue;

      /* We've found a heap with (potentially) enough space */
      if ((ret = per_heap_kmalloc(&pos->heap, size, flags))) {

         /*
          * The alloc succeeded. If now the heap is full, remove it from the
          * 'avail' list.
          */

         if (pos->heap.mem_allocated == pos->heap.size)
            remove_from_avail_list(pos);

         return ret;
      }
   }

   if (!(new_node = alloc_new_small_heap()))
      return NULL;
{'fraction_non_alphanumeric': 0.10151783954267692, 'fraction_numerical': 0.00384387936132466, 'mean_word_length': 3.347472150814053, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 67, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6519217', 'n_tokens_mistral': 3656, 'n_tokens_neox': 3506, 'n_words': 701}