Re: [PATCH RFC v2 06/10] slab: sheaf prefilling for guaranteed allocations

[Vlastimil Babka <vbabka@suse.cz>]

On 2/23/25 04:54, Suren Baghdasaryan wrote:
> On Fri, Feb 14, 2025 at 8:27 AM Vlastimil Babka <vbabka@suse.cz> wrote:
>>
>> Add functions for efficient guaranteed allocations e.g. in a critical
>> section that cannot sleep, when the exact number of allocations is not
>> known beforehand, but an upper limit can be calculated.
>>
>> kmem_cache_prefill_sheaf() returns a sheaf containing at least given
>> number of objects.
>>
>> kmem_cache_alloc_from_sheaf() will allocate an object from the sheaf
>> and is guaranteed not to fail until depleted.
>>
>> kmem_cache_return_sheaf() is for giving the sheaf back to the slab
>> allocator after the critical section. This will also attempt to refill
>> it to cache's sheaf capacity for better efficiency of sheaves handling,
>> but it's not stricly necessary to succeed.
>>
>> kmem_cache_refill_sheaf() can be used to refill a previously obtained
>> sheaf to requested size. If the current size is sufficient, it does
>> nothing. If the requested size exceeds cache's sheaf_capacity and the
>> sheaf's current capacity, the sheaf will be replaced with a new one,
>> hence the indirect pointer parameter.
>>
>> kmem_cache_sheaf_size() can be used to query the current size.
>>
>> The implementation supports requesting sizes that exceed cache's
>> sheaf_capacity, but it is not efficient - such sheaves are allocated
>> fresh in kmem_cache_prefill_sheaf() and flushed and freed immediately by
>> kmem_cache_return_sheaf(). kmem_cache_refill_sheaf() might be expecially
> 
> s/expecially/especially
> 
>> ineffective when replacing a sheaf with a new one of a larger capacity.
>> It is therefore better to size cache's sheaf_capacity accordingly.
> 
> If support for sizes exceeding sheaf_capacity adds much complexity
> with no performance benefits, I think it would be ok not to support
> them at all. Users know the capacity of a particular kmem_cache, so
> they can use this API only when their needs are within sheaf_capacity,
> otherwise either size the sheaf appropriately or use slab bulk
> allocation.

As Harry explained, the users (e.g. maple tree) would have to implement the
fallback for unusual situations instead, so it's better to implement it just
once here.

>>
>> Signed-off-by: Vlastimil Babka <vbabka@suse.cz>
> 
> Reviewed-by: Suren Baghdasaryan <surenb@google.com>

Thanks.

>> +/*
>> + * Use this to return a sheaf obtained by kmem_cache_prefill_sheaf()
>> + * It tries to refill the sheaf back to the cache's sheaf_capacity
>> + * to avoid handling partially full sheaves.
>> + *
>> + * If the refill fails because gfp is e.g. GFP_NOWAIT, the sheaf is
>> + * instead dissolved
> 
> Refilling the sheaf here assumes that in the future we are more likely
> to allocate than to free objects or shrink the slab. If the reverse is
> true then it would make sense to flush the sheaf and add it as an
> empty one into the barn. The fact that flushing can't fail would be
> another advantage... We don't know the future but should we be
> predicting a more costly case?

What the comment doesn't say is we first try to make the sheaf become
pcs->spare without any refill. This is the ideal scenario if nobody
interrupts us between prefill (we grab the spare) and return (we return the
spare).

Also the refill is only attempted if the barn can accept full sheaf.

I have clarified the comment.

Maybe we could make the decision to flush e.g. if the sheaf is below half of
the capacity, but that can be subject to further performance evaluation.

>> + */
>> +void kmem_cache_return_sheaf(struct kmem_cache *s, gfp_t gfp,
>> +                            struct slab_sheaf *sheaf)
>> +{
>> +       struct slub_percpu_sheaves *pcs;
>> +       bool refill = false;
>> +       struct node_barn *barn;
>> +
>> +       if (unlikely(sheaf->capacity != s->sheaf_capacity)) {
>> +               sheaf_flush(s, sheaf);
>> +               kfree(sheaf);
>> +               return;
>> +       }
>> +
>> +       localtry_lock(&s->cpu_sheaves->lock);
>> +       pcs = this_cpu_ptr(s->cpu_sheaves);
>> +
>> +       if (!pcs->spare) {
>> +               pcs->spare = sheaf;
>> +               sheaf = NULL;
>> +       } else if (pcs->barn->nr_full >= MAX_FULL_SHEAVES) {
>> +               /* racy check */
>> +               barn = pcs->barn;
>> +               refill = true;
>> +       }
>> +
>> +       localtry_unlock(&s->cpu_sheaves->lock);
>> +
>> +       if (!sheaf)
>> +               return;
>> +
>> +       /*
>> +        * if the barn is full of full sheaves or we fail to refill the sheaf,
>> +        * simply flush and free it
>> +        */
>> +       if (!refill || refill_sheaf(s, sheaf, gfp)) {
>> +               sheaf_flush(s, sheaf);
>> +               free_empty_sheaf(s, sheaf);
>> +               return;
>> +       }
>> +
>> +       /* we racily determined the sheaf would fit, so now force it */
>> +       barn_put_full_sheaf(barn, sheaf, true);
>> +}
>> +
>> +/*
>> + * refill a sheaf previously returned by kmem_cache_prefill_sheaf to at least
>> + * the given size
>> + *
>> + * the sheaf might be replaced by a new one when requesting more than
>> + * s->sheaf_capacity objects if such replacement is necessary, but the refill
>> + * fails (with -ENOMEM), the existing sheaf is left intact
>> + */
>> +int kmem_cache_refill_sheaf(struct kmem_cache *s, gfp_t gfp,
>> +                           struct slab_sheaf **sheafp, unsigned int size)
>> +{
>> +       struct slab_sheaf *sheaf;
>> +
>> +       /*
>> +        * TODO: do we want to support *sheaf == NULL to be equivalent of
>> +        * kmem_cache_prefill_sheaf() ?
>> +        */
>> +       if (!sheafp || !(*sheafp))
>> +               return -EINVAL;
>> +
>> +       sheaf = *sheafp;
>> +       if (sheaf->size >= size)
>> +               return 0;
>> +
>> +       if (likely(sheaf->capacity >= size)) {
>> +               if (likely(sheaf->capacity == s->sheaf_capacity))
>> +                       return refill_sheaf(s, sheaf, gfp);
>> +
>> +               if (!__kmem_cache_alloc_bulk(s, gfp, sheaf->capacity - sheaf->size,
>> +                                            &sheaf->objects[sheaf->size])) {
>> +                       return -ENOMEM;
>> +               }
>> +               sheaf->size = sheaf->capacity;
>> +
>> +               return 0;
>> +       }
>> +
>> +       /*
>> +        * We had a regular sized sheaf and need an oversize one, or we had an
>> +        * oversize one already but need a larger one now.
>> +        * This should be a very rare path so let's not complicate it.
>> +        */
>> +       sheaf = kmem_cache_prefill_sheaf(s, gfp, size);
> 
> WIth all the above I think you always end up refilling up to
> sheaf->capacity. Not sure if we should mention that in the comment for
> this function because your statement about refilling to at least the
> given size is still correct.

OK mentioned it in the comment.