pctrie: Add a function for calculating the number of allocations for inserting a range of keys

[bnovkov]

Summary

This patch adds a function that uses a pctrie_iterator to determine the exact number of pctrie nodes required for inserting a range of keys.
First, the function calculates the number of nodes required to insert the key range into an empty tree, and then uses the iterator to look up the first and last key
and decrement the node count if any shared nodes exist.

The function turned out more complex than I'd anticipated, so I'll break things down into several smaller parts and do my best to describe each one.
First off, we most straightforward case—inserting keys from [0, N).
For a complete pctrie (i.e., where N is a power of PCTRIE_WIDTH), the total number of required nodes can be calculated using N / (PCTRIE_COUNT - 1).
This is relatively straightforward to prove—since the tree is full, it will have N/(PCTRIE_COUNT) level 0 nodes, N/(PCTRIE_COUNT^2) level 1 nodes, and so on.
This is a convergent geometric series with a = N, r = 1/ PCTRIE_COUNT, so this whole expression can be further simplified as follows:

After substituting r with 1/ PCTRIE_COUNT and simplifying, the final expression is N / (PCTRIE_COUNT - 1).

Next, we have to account for path compression. Here I used the fact that each pctrie can be broken down into a "leftmost" and "rightmost" incomplete pctrie, with 0 or multiple complete pctries in between.
This algorithm in pctrie_subtree_nnodes properly accounts for path compression in the "rightmost" subtree, and works by splitting the pctrie into multiple smaller subtrees, adding a top-level node for each incomplete subtree.
For instance, given PCTRIE_COUNT = 16 and N = 288, the algorithm will split the tree into 256 / 15 + (32 / 16 + 1) + 1, giving us 21 nodes in total. Path compression in the "leftmost" subtree is calculated with the same algorithm, using the [0, lkeys) range, where lkeys is the total number of keys covered by the "leftmost" subtree.

After calculating the number of nodes using the previously described algorithm, the function will look up the first and last key and adjust the node count accordingly.
The node count will be incremented if a value is stored in a target slot. The function will then inspect every node in the iterator's path and decrement the node count to avoid allocating any nodes that might be present in the tree.
I've left comments with additional explanations before each node count adjustment.

Path compression introduces plenty of edge cases that had to be addressed, which is why the function is currently a bit verbose. The logic could probably be simplified even further, but I'll get back to this after a short break. I'll stay on top of any feedback you might have meanwhile.

Test Plan

This patch was evaluated using two tests.

The first test consisted of a kernel module inserting keys from [0, N) for increasingly larger values of N and was meant to test pctrie_subtree_nnodes.
The second test involved plugging the patch into pctrie_batch_insert and testing it using the batch-aware vm_page_grab_pages implemented in D47051.
The setup involved a UFS VM running a ktls-enabled nginx that was stress-tested using ApacheBench.

The first test completes without any errors, and so far I've had no cases where the node count was overestimated or underestimated while running the second test.

[dougm]

sys/kern/subr_pctrie.c

691

Wrong for end = 1 mod PCTRIE_COUNT.

Here's a simpler, and correct, implementation.

	int cnt = 0;
	end--;
	while (end > 0) {
		if (end % PCTRIE_COUNT != 0)
			++cnt;
		end /= PCTRIE_COUNT;
		cnt += end;
	}
	return (cnt);

If I had a version without the loop, I'd offer it, but I don't at the moment.