homa_pool.c

// SPDX-License-Identifier: BSD-2-Clause OR GPL-2.0+

#include "homa_impl.h"
#ifndef __STRIP__ /* See strip.py */
#include "homa_grant.h"
#endif /* See strip.py */
#include "homa_pool.h"

/* This file contains functions that manage user-space buffer pools. */

/* Used when determining how many bpages to consider for allocation. */
#define MIN_EXTRA 4

/**
 * set_bpages_needed() - Set the bpages_needed field of @pool based
 * on the length of the first RPC that's waiting for buffer space.
 * @pool: Pool to update. Must be locked by caller.
 */
static void set_bpages_needed(struct homa_pool *pool)
	__must_hold(pool->lock)
{
	struct homa_rpc *rpc = list_first_entry(&pool->waiting_for_bufs,
						struct homa_rpc, buf_links);

	pool->bpages_needed = (rpc->msgin.length + HOMA_BPAGE_SIZE - 1) >>
			      HOMA_BPAGE_SHIFT;
}

/**
 * homa_pool_alloc() - Allocate and initialize a new homa_pool (it will have
 * no region associated with it until homa_pool_set_region is invoked).
 * @hsk:          Socket the pool will be associated with.
 * Return: A pointer to the new pool or a negative errno.
 */
struct homa_pool *homa_pool_alloc(struct homa_sock *hsk)
{
	struct homa_pool *pool;

	pool = kzalloc(sizeof(*pool), GFP_KERNEL_ACCOUNT);
	if (!pool)
		return ERR_PTR(-ENOMEM);
	spin_lock_init(&pool->lock);
	INIT_LIST_HEAD(&pool->waiting_for_bufs);
	pool->hsk = hsk;
	return pool;
}

/**
 * homa_pool_set_region() - Associate a region of memory with a pool.
 * @pool:         Buffer pool the region will be associated with. Must not
 *                currently have a region associated with it.
 * @region:       First byte of the memory region for the pool, allocated
 *                by the application; must be page-aligned.
 * @region_size:  Total number of bytes available at @buf_region.
 * Return: Either zero (for success) or a negative errno for failure.
 */
int homa_pool_set_region(struct homa_pool *pool, void __user *region,
			 u64 region_size)
{
	struct homa_pool_core __percpu *cores;
	struct homa_bpage *descriptors;
	int i, result, num_bpages;
	u64 min_size;

	if (((uintptr_t)region) & ~PAGE_MASK) {
		pool->hsk->error_msg = "buffer pool is not page aligned";
		return -EINVAL;
	}

	/* The region must be large enough to hold a private bpage for
	 * each core, plus enough additional space for a couple of
	 * full-size messages. Otherwise we could deadlock over allocation
	 * (no free pages available, so no attempt is made to steal back
	 * private pages).
	 */
	min_size = nr_cpu_ids;
	min_size = min_size * HOMA_BPAGE_SIZE + 2 * HOMA_MAX_MESSAGE_LENGTH;
	if (region_size < min_size) {
		pool->hsk->error_msg = "buffer pool is not large enough";
		return -EINVAL;
	}

	/* The region must not be so large that offsets into the region
	 * overflow the 32-bit values used in bpage_offsets in
	 * homa_recvmsg_args.
	 */
	if (region_size > 0x100000000ULL) {
		pool->hsk->error_msg = "buffer pool cannot be larger than 4 GB";
		return -EINVAL;
	}

	/* Allocate memory before locking the pool, so we can allocate
	 * without GFP_ATOMIC.
	 */
	num_bpages = region_size >> HOMA_BPAGE_SHIFT;
	descriptors = kmalloc_array(num_bpages, sizeof(struct homa_bpage),
				    GFP_KERNEL | __GFP_ZERO | __GFP_ACCOUNT |
				    __GFP_NOWARN);
	if (!descriptors)
		return -ENOMEM;
	cores = alloc_percpu_gfp(struct homa_pool_core,
				 GFP_KERNEL_ACCOUNT | __GFP_ZERO);
	if (!cores) {
		result = -ENOMEM;
		goto error;
	}

	homa_pool_lock(pool);
	if (pool->region) {
		result = -EINVAL;
		homa_pool_unlock(pool);
		goto error;
	}

	pool->num_bpages = num_bpages;
	pool->descriptors = descriptors;
	atomic_set(&pool->free_bpages, pool->num_bpages);
	pool->bpages_needed = INT_MAX;
	pool->cores = cores;
	pool->check_waiting_invoked = 0;

	for (i = 0; i < pool->num_bpages; i++) {
		struct homa_bpage *bp = &pool->descriptors[i];

		spin_lock_init(&bp->lock);
		bp->owner = -1;
	}

	/* Barrier ensures that concurrent cores trying to access
	 * newly-created region will see all the region's data.
	 */
	smp_store_release(&pool->region, (char __user *)region);
	homa_pool_unlock(pool);
	return 0;

error:
	kfree(descriptors);
	free_percpu(cores);
	return result;
}

/**
 * homa_pool_free() - Destructor for homa_pool. After this method
 * returns, the object should not be used (it will be freed here).
 * @pool: Pool to destroy.
 */
void homa_pool_free(struct homa_pool *pool)
{
#ifdef __UNIT_TEST__
	mock_free_pool(pool);
#endif /* __UNIT_TEST__ */
	if (homa_pool_exists(pool)) {
		kfree(pool->descriptors);
		free_percpu(pool->cores);
		pool->region = NULL;
	}
	kfree(pool);
}

/**
 * homa_pool_get_rcvbuf() - Return information needed to handle getsockopt
 * for HOMA_SO_RCVBUF.
 * @pool:         Pool for which information is needed.
 * @args:         Store info here.
 */
void homa_pool_get_rcvbuf(struct homa_pool *pool,
			  struct homa_rcvbuf_args *args)
{
	args->start = (uintptr_t)pool->region;
	args->length = (u64)pool->num_bpages << HOMA_BPAGE_SHIFT;
}

/**
 * homa_bpage_available() - Check whether a bpage is available for use.
 * @bpage:      Bpage to check
 * @now:        Current time (homa_clock() units)
 * Return:      True if the bpage is free or if it can be stolen, otherwise
 *              false.
 */
bool homa_bpage_available(struct homa_bpage *bpage, u64 now)
{
	return bpage->refs == 0 || (bpage->refs == 1 && bpage->owner >= 0 &&
				    bpage->expiration <= now);
}

/**
 * homa_pool_get_pages() - Allocate one or more full pages from the pool.
 * @pool:         Pool from which to allocate pages
 * @num_pages:    Number of pages needed
 * @pages:        The indices of the allocated pages are stored here; caller
 *                must ensure this array is big enough. Reference counts have
 *                been set to 1 on all of these pages (or 2 if set_owner
 *                was specified).
 * @set_owner:    If nonzero, the current core is marked as owner of all
 *                of the allocated pages (and the expiration time is also
 *                set). Otherwise the pages are left unowned.
 * Return: 0 for success, -1 if there wasn't enough free space in the pool.
 */
int homa_pool_get_pages(struct homa_pool *pool, int num_pages, u32 *pages,
			int set_owner)
{
	int core_num = smp_processor_id();
	struct homa_pool_core *core;
	u64 now = homa_clock();
	int alloced = 0;
	int limit = 0;
	int free;

	core = this_cpu_ptr(pool->cores);

	/* Atomically check for available space and decrement the
	 * free page counter.
	 */
	free = atomic_read_acquire(&pool->free_bpages);
	while (1) {
		if (free < num_pages)
			return -1;
		if (atomic_try_cmpxchg(&pool->free_bpages, &free,
				       free - num_pages))
			break;
	}

	/* Once we get to this point we know we will be able to find
	 * enough free pages; now we just have to find them.
	 */
	while (alloced != num_pages) {
		struct homa_bpage *bpage;
		int cur;

		/* If we don't need to use all of the bpages in the pool,
		 * then try to use only the ones with low indexes. This
		 * will reduce the cache footprint for the pool by reusing
		 * a few bpages over and over. Specifically this code will
		 * not consider any candidate page whose index is >= limit.
		 * Limit is chosen to make sure there are a reasonable
		 * number of free pages in the range, so we won't have to
		 * check a huge number of pages.
		 */
		if (limit == 0) {
			int extra;

			limit = pool->num_bpages -
				atomic_read(&pool->free_bpages);
			extra = limit >> 2;
			limit += (extra < MIN_EXTRA) ? MIN_EXTRA : extra;
			if (limit > pool->num_bpages)
				limit = pool->num_bpages;
		}

		cur = core->next_candidate;
		core->next_candidate++;
		if (cur >= limit) {
			core->next_candidate = 0;

			/* Must recompute the limit for each new loop through
			 * the bpage array: we may need to consider a larger
			 * range of pages because of concurrent allocations.
			 */
			limit = 0;
			continue;
		}
		bpage = &pool->descriptors[cur];

		/* Figure out whether this candidate is free (or can be
		 * stolen). Do a quick check without locking the page, and
		 * if the page looks promising, then lock it and check again
		 * (must check again in case someone else snuck in and
		 * grabbed the page).
		 */
		if (!homa_bpage_available(bpage, now))
			continue;
		if (!spin_trylock_bh(&bpage->lock))
			/* Rather than wait for a locked page to become free,
			 * just go on to the next page. If the page is locked,
			 * it probably won't turn out to be available anyway.
			 */
			continue;
		if (!homa_bpage_available(bpage, now)) {
			spin_unlock_bh(&bpage->lock);
			continue;
		}
		if (bpage->owner >= 0)
			atomic_fetch_inc_release(&pool->free_bpages);
		if (set_owner) {
			bpage->refs = 2;
			bpage->owner = core_num;
			bpage->expiration = now +
					    pool->hsk->homa->bpage_lease_cycles;
		} else {
			bpage->refs = 1;
			bpage->owner = -1;
		}
		spin_unlock_bh(&bpage->lock);
		pages[alloced] = cur;
		alloced++;
	}
	return 0;
}

/**
 * homa_pool_alloc_msg() - Allocate buffer space for an incoming message.
 * @rpc:  RPC that needs space allocated for its incoming message (space must
 *        not already have been allocated). The fields @msgin->num_buffers
 *        and @msgin->buffers are filled in. Must be locked by caller.
 * Return: The return value is normally 0, which means either buffer space
 * was allocated or the @rpc was queued on @hsk->waiting. If a fatal error
 * occurred, such as no buffer pool present, then a negative errno is
 * returned.
 */
int homa_pool_alloc_msg(struct homa_rpc *rpc)
	__must_hold(rpc->bucket->lock)
{
	struct homa_pool *pool = rpc->hsk->buffer_pool;
	int full_pages, partial, i, core_id;
	struct homa_pool_core *core;
	u32 pages[HOMA_MAX_BPAGES];
	struct homa_bpage *bpage;
	struct homa_rpc *other;

	if (!homa_pool_exists(pool))
		return -ENOMEM;
	if (rpc->state == RPC_DEAD)
		return 0;

	/* First allocate any full bpages that are needed. */
	full_pages = rpc->msgin.length >> HOMA_BPAGE_SHIFT;
	if (unlikely(full_pages)) {
		if (homa_pool_get_pages(pool, full_pages, pages, 0) != 0)
			goto out_of_space;
		for (i = 0; i < full_pages; i++)
			rpc->msgin.bpage_offsets[i] = pages[i] <<
					HOMA_BPAGE_SHIFT;
	}
	rpc->msgin.num_bpages = full_pages;

	/* The last chunk may be less than a full bpage; for this we use
	 * the bpage owned by this core (and reuse it for multiple messages).
	 */
	partial = rpc->msgin.length & (HOMA_BPAGE_SIZE - 1);
	if (unlikely(partial == 0))
		goto success;
	core_id = smp_processor_id();
	core = this_cpu_ptr(pool->cores);
	bpage = &pool->descriptors[core->page_hint];
	spin_lock_bh(&bpage->lock);
	if (bpage->owner != core_id) {
		spin_unlock_bh(&bpage->lock);
		goto new_page;
	}
	if ((core->allocated + partial) > HOMA_BPAGE_SIZE) {
#ifndef __STRIP__ /* See strip.py */
		if (bpage->refs == 1) {
			/* Bpage is totally free, so we can reuse it. */
			core->allocated = 0;
			INC_METRIC(bpage_reuses, 1);
#else /* See strip.py */
		if (bpage->refs == 1) {
			/* Bpage is totally free, so we can reuse it. */
			core->allocated = 0;
#endif /* See strip.py */
		} else {
			bpage->owner = -1;

			/* We know the reference count can't reach zero here
			 * because of check above, so we won't have to decrement
			 * pool->free_bpages.
			 */
			bpage->refs--;
			spin_unlock_bh(&bpage->lock);
			goto new_page;
		}
	}
	bpage->expiration = homa_clock() +
			    pool->hsk->homa->bpage_lease_cycles;
	bpage->refs++;
	spin_unlock_bh(&bpage->lock);
	goto allocate_partial;

	/* Can't use the current page; get another one. */
new_page:
	if (homa_pool_get_pages(pool, 1, pages, 1) != 0) {
		homa_pool_free_bufs(pool, rpc->msgin.num_bpages,
				    rpc->msgin.bpage_offsets);
		rpc->msgin.num_bpages = 0;
		goto out_of_space;
	}
	core->page_hint = pages[0];
	core->allocated = 0;

allocate_partial:
	rpc->msgin.bpage_offsets[rpc->msgin.num_bpages] = core->allocated
			+ (core->page_hint << HOMA_BPAGE_SHIFT);
	rpc->msgin.num_bpages++;
	core->allocated += partial;

success:
	tt_record4("Allocated %d bpage pointers on port %d for id %d, free_bpages now %d",
		   rpc->msgin.num_bpages, pool->hsk->port, rpc->id,
		   atomic_read(&pool->free_bpages));
	return 0;

	/* We get here if there wasn't enough buffer space for this
	 * message; add the RPC to pool->waiting_for_bufs. The list is sorted
	 * by RPC length in order to implement SRPT.
	 */
out_of_space:
	INC_METRIC(buffer_alloc_failures, 1);
	tt_record4("Buffer allocation failed, port %d, id %d, length %d, free_bpages %d",
		   pool->hsk->port, rpc->id, rpc->msgin.length,
		   atomic_read(&pool->free_bpages));
	homa_pool_lock(pool);
	list_for_each_entry(other, &pool->waiting_for_bufs, buf_links) {
		if (other->msgin.length > rpc->msgin.length) {
			list_add_tail(&rpc->buf_links, &other->buf_links);
			goto queued;
		}
	}
	list_add_tail(&rpc->buf_links, &pool->waiting_for_bufs);

queued:
	set_bpages_needed(pool);
	homa_pool_unlock(pool);
	return 0;
}

/**
 * homa_pool_get_buffer() - Given an RPC, figure out where to store incoming
 * message data.
 * @rpc:        RPC for which incoming message data is being processed; its
 *              msgin must be properly initialized and buffer space must have
 *              been allocated for the message.
 * @offset:     Offset within @rpc's incoming message.
 * @available:  Will be filled in with the number of bytes of space available
 *              at the returned address (could be zero if offset is
 *              (erroneously) past the end of the message).
 * Return:      The application's virtual address for buffer space corresponding
 *              to @offset in the incoming message for @rpc.
 */
void __user *homa_pool_get_buffer(struct homa_rpc *rpc, int offset,
				  int *available)
{
	int bpage_index, bpage_offset;

	bpage_index = offset >> HOMA_BPAGE_SHIFT;
	if (offset >= rpc->msgin.length) {
		WARN_ONCE(true, "%s got offset %d >= message length %d\n",
			  __func__, offset, rpc->msgin.length);
		*available = 0;
		return NULL;
	}
	bpage_offset = offset & (HOMA_BPAGE_SIZE - 1);
	*available = (bpage_index < (rpc->msgin.num_bpages - 1))
			? HOMA_BPAGE_SIZE - bpage_offset
			: rpc->msgin.length - offset;
	return rpc->hsk->buffer_pool->region +
			rpc->msgin.bpage_offsets[bpage_index] + bpage_offset;
}

/**
 * homa_pool_free_bufs() - Release buffer space so that it can be reused.
 * @pool:         Pool that the buffer space belongs to. Doesn't need to
 *                be locked.
 * @num_buffers:  How many buffers to release.
 * @buffers:      Points to @num_buffers values, each of which is an offset
 *                from the start of the pool to the buffer to be released.
 * Return:        0 for success, otherwise a negative errno.
 */
int homa_pool_free_bufs(struct homa_pool *pool, int num_buffers, u32 *buffers)
{
	int i;

	/* Note: a misbehaving application could potentially free the same
	 * offset multiple times, which could lead to accidental sharing
	 * of buffer space across RPCs. Fortunately, the resulting corruption
	 * affects only the misbehaving application (buffer pools are not
	 * shared across applciations), and eliminating any potential for
	 * corruption seems hard.
	 */

	if (!homa_pool_exists(pool))
		return -EINVAL;
	for (i = 0; i < num_buffers; i++) {
		u32 bpage_index = buffers[i] >> HOMA_BPAGE_SHIFT;
		struct homa_bpage *bpage;

		if (bpage_index >= pool->num_bpages)
			return -EINVAL;
		bpage = &pool->descriptors[bpage_index];

		spin_lock_bh(&bpage->lock);
		if (bpage->refs == 0 || (bpage->refs == 1 &&
					 bpage->owner >= 0)) {
			/* This bpage is already free; looks like the app
			 * has misbehaved and freed an offset multiple times.
			 */
			spin_unlock_bh(&bpage->lock);
			return -EINVAL;
		}
		bpage->refs--;
		if (bpage->refs == 0)
			atomic_fetch_inc(&pool->free_bpages);
		spin_unlock_bh(&bpage->lock);
	}
	tt_record3("Released %d bpages, free_bpages for port %d now %d",
		   num_buffers, pool->hsk->port,
		   atomic_read(&pool->free_bpages));
	return 0;
}

/**
 * homa_pool_check_waiting() - Checks to see if there are enough free
 * bpages to wake up any RPCs that were blocked. Whenever
 * homa_pool_free_bufs is invoked, this function must be invoked later,
 * at a point when the caller holds no locks (homa_pool_free_bufs may
 * be invoked with locks held, so it can't safely invoke this function).
 * This is regrettably tricky, but I can't think of a better solution.
 * @pool:         Information about the buffer pool.
 */
void homa_pool_check_waiting(struct homa_pool *pool)
{
#ifdef __UNIT_TEST__
	pool->check_waiting_invoked += 1;
#endif /* __UNIT_TEST__ */
	if (!homa_pool_exists(pool))
		return;
	while (atomic_read_acquire(&pool->free_bpages) >= pool->bpages_needed) {
		struct homa_rpc *rpc;

		homa_pool_lock(pool);
		if (list_empty(&pool->waiting_for_bufs)) {
			pool->bpages_needed = INT_MAX;
			homa_pool_unlock(pool);
			break;
		}
		rpc = list_first_entry(&pool->waiting_for_bufs,
				       struct homa_rpc, buf_links);
		if (!homa_rpc_try_lock(rpc)) {
			/* Can't spin on the RPC lock because we're holding
			 * the pool lock and the lock order is rpc-then-pool
			 * (see "Homa Locking Strategy" in homa_impl.h).
			 * Instead, release the pool lock and try the entire
			 * operation again.
			 */
			homa_pool_unlock(pool);
			UNIT_LOG("; ", "rpc lock unavailable in %s", __func__);
			cpu_relax();
			continue;
		}
		list_del_init(&rpc->buf_links);
		if (list_empty(&pool->waiting_for_bufs))
			pool->bpages_needed = INT_MAX;
		else
			set_bpages_needed(pool);
		homa_pool_unlock(pool);
		tt_record4("Retrying buffer allocation for id %d, length %d, free_bpages %d, new bpages_needed %d",
			   rpc->id, rpc->msgin.length,
			   atomic_read(&pool->free_bpages),
			   pool->bpages_needed);
		homa_pool_alloc_msg(rpc);
		if (rpc->msgin.num_bpages > 0)
			homa_pool_wakeup_rpc(rpc);
		homa_rpc_unlock(rpc);
	}
}

/**
 * homa_pool_wakeup_rpc() - This function is invoked when an RPC that had
 * blocked waiting for buffer space finally gets the space it needs. It
 * arranges for transmission of the RPC's data to resume.
 * @rpc:    RPC to wake up
 */
void homa_pool_wakeup_rpc(struct homa_rpc *rpc)
	__must_hold(rpc->bucket->lock)
{
	struct homa_resend_hdr resend;

#ifndef __STRIP__ /* See strip.py */
	/* If the RPC is scheduled, all we have to do is issue a grant. */
	if (test_bit(RPC_GRANTABLE, &rpc->flags)) {
		homa_grant_check_rpc(rpc);
		return;
	}
#endif /* See strip.py */

	/* Unscheduled: must ask for the message to be retransmitted. */
	resend.offset = htonl(0);
	resend.length = htonl(rpc->msgin.length);

#ifndef __STRIP__ /* See strip.py */
	/* Use the next-to-highest priority level to provide a priority boost
	 * without interfering with the highest priority traffic such as
	 * control packets.
	 */
	resend.priority = homa_high_priority(rpc->hsk->homa);
#endif /* See strip.py */
	homa_xmit_control(RESEND, &resend, sizeof(resend), rpc);
}

/**
 * homa_pool_avail_bytes() - Return a count of the number of bytes currently
 * unused and available for allocation in a pool.
 * @pool:    Pool of interest.
 * Return:    See above.
 */
u64 homa_pool_avail_bytes(struct homa_pool *pool)
{
	struct homa_pool_core *core;
	struct homa_bpage *bpage;
	u64 avail;
	int cpu;

	if (!homa_pool_exists(pool))
		return 0;
	avail = atomic_read(&pool->free_bpages);
	avail *= HOMA_BPAGE_SIZE;
	for (cpu = 0; cpu < nr_cpu_ids; cpu++) {
		core = per_cpu_ptr(pool->cores, cpu);
		bpage = &pool->descriptors[core->page_hint];
		if (bpage->owner == cpu) {
			if (bpage->refs > 1)
				avail += HOMA_BPAGE_SIZE - core->allocated;
			else
				avail += HOMA_BPAGE_SIZE;
		}
	}
	return avail;
}

#ifndef __STRIP__ /* See strip.py */
/**
 * homa_pool_lock_slow() - This function implements the slow path for
 * acquiring the lock for a buffer pool. It is invoked when the lock isn't
 * immediately available. It waits for the lock but also records statistics
 * about the waiting time.
 * @pool:    Buffer pool to  lock.
 */
void homa_pool_lock_slow(struct homa_pool *pool)
	__acquires(pool->lock)
{
	u64 start = homa_clock();

	tt_record("beginning wait for bpool lock");
	spin_lock_bh(&pool->lock);
	tt_record("ending wait for bpool lock");
	INC_METRIC(bpool_lock_misses, 1);
	INC_METRIC(bpool_lock_miss_cycles, homa_clock() - start);
}
#endif /* See strip.py */