Skip to content

Commit

Permalink
map: Introduce BatchCursor abstraction
Browse files Browse the repository at this point in the history 
The batch API in the kernel is quite tricky to get right. In order to
pass the arguments in correctly, especially regarding the "in_batch" and
"out_batch" arguments, it requires knowledge of how the batch map
iteration is implemented in the kernel.

In order to avoid pushing the cognitive overhead onto users of the
library, abstract away the details of batch map iteration into an opaque
type, BatchCursor.

Users are expected to pass in a reference to a BatchCursor and the
library will handle allocating the underlying buffer when needed.
Specifically, the first time that a batch API is called using the
cursor, the underlying buffer will be nil, which signfifies that the
"in_batch" (aka "prevKey") should be nil to indicate the starting of a
batching operation, and "out_batch" (aka "nextKey") is set to an
newly-allocated buffer. Previously, users were expected to carry the
"out_batch" reference in order to pass it in as "in_batch" upon the next
call to the batching operation. Instead, with this abstraction, the
cursor handles it for the user, by using the underlying buffer for both
"in_batch" and "out_batch" in subsequent calls.

Signed-off-by: Chris Tarazi <chris@isovalent.com>
Suggested-by: Lorenz Bauer <lmb@isovalent.com>
  • Loading branch information
@christarazi @lmb
christarazi authored and lmb committed Dec 1, 2023
1 parent f0d238d commit d9f8904
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 41 deletions.
56 changes: 32 additions & 24 deletions map.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"io"
"math"
"math/rand"
"os"
"path/filepath"
Expand Down Expand Up @@ -963,32 +964,51 @@ func (m *Map) guessNonExistentKey() ([]byte, error) {
//
// "keysOut" and "valuesOut" must be of type slice, a pointer
// to a slice or buffer will not work.
// "prevKey" is the key to start the batch lookup from, it will
// *not* be included in the results. Use nil to start at the first key.
// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
// "cursor" to subsequent calls of this function to continue the batching
// operation in the case of chunking.
//
// ErrKeyNotExist is returned when the batch lookup has reached
// the end of all possible results, even when partial results
// are returned. It should be used to evaluate when lookup is "done".
func (m *Map) BatchLookup(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
func (m *Map) BatchLookup(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
return m.batchLookup(sys.BPF_MAP_LOOKUP_BATCH, cursor, keysOut, valuesOut, opts)
}

// BatchLookupAndDelete looks up many elements in a map at once,
//
// It then deletes all those elements.
// "keysOut" and "valuesOut" must be of type slice, a pointer
// to a slice or buffer will not work.
// "prevKey" is the key to start the batch lookup from, it will
// *not* be included in the results. Use nil to start at the first key.
// "cursor" is an pointer to an opaque handle. It must be non-nil. Pass
// "cursor" to subsequent calls of this function to continue the batching
// operation in the case of chunking.
//
// ErrKeyNotExist is returned when the batch lookup has reached
// the end of all possible results, even when partial results
// are returned. It should be used to evaluate when lookup is "done".
func (m *Map) BatchLookupAndDelete(prevKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, prevKey, nextKeyOut, keysOut, valuesOut, opts)
func (m *Map) BatchLookupAndDelete(cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
return m.batchLookup(sys.BPF_MAP_LOOKUP_AND_DELETE_BATCH, cursor, keysOut, valuesOut, opts)
}

func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
// BatchCursor represents a starting point for a batch operation.
type BatchCursor struct {
opaque []byte // len() is max(key_size, 4) to be on the safe side
}

func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut interface{}, opts *BatchOptions) (int, error) {
var inBatch []byte
if cursor.opaque == nil {
// * generic_map_lookup_batch requires that batch_out is key_size bytes.
// This is used by array and LPM maps.
//
// * __htab_map_lookup_and_delete_batch requires u32. This is used by the
// various hash maps.
cursor.opaque = make([]byte, int(math.Max(float64(m.keySize), 4)))
} else {
inBatch = cursor.opaque
}

if err := haveBatchAPI(); err != nil {
return 0, err
}
Expand All @@ -1011,40 +1031,28 @@ func (m *Map) batchLookup(cmd sys.Cmd, startKey, nextKeyOut, keysOut, valuesOut
keyPtr := sys.NewSlicePointer(keyBuf)
valueBuf := make([]byte, count*int(m.fullValueSize))
valuePtr := sys.NewSlicePointer(valueBuf)
nextBuf := makeMapSyscallOutput(nextKeyOut, int(m.keySize))

attr := sys.MapLookupBatchAttr{
MapFd: m.fd.Uint(),
Keys: keyPtr,
Values: valuePtr,
Count: uint32(count),
OutBatch: nextBuf.Pointer(),
InBatch: sys.NewSlicePointer(inBatch),
OutBatch: sys.NewSlicePointer(cursor.opaque),
}

if opts != nil {
attr.ElemFlags = opts.ElemFlags
attr.Flags = opts.Flags
}

var err error
if startKey != nil {
attr.InBatch, err = marshalMapSyscallInput(startKey, int(m.keySize))
if err != nil {
return 0, err
}
}

_, sysErr := sys.BPF(cmd, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
sysErr = wrapMapError(sysErr)
if sysErr != nil && !errors.Is(sysErr, unix.ENOENT) {
return 0, sysErr
}

err = nextBuf.Unmarshal(nextKeyOut)
if err != nil {
return 0, err
}
err = sysenc.Unmarshal(keysOut, keyBuf)
err := sysenc.Unmarshal(keysOut, keyBuf)
if err != nil {
return 0, err
}
Expand Down
33 changes: 16 additions & 17 deletions map_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,6 @@ func TestBatchAPIArray(t *testing.T) {
defer m.Close()

var (
nextKey uint32
keys = []uint32{0, 1}
values = []uint32{42, 4242}
lookupKeys = make([]uint32, 2)
Expand All @@ -133,24 +132,23 @@ func TestBatchAPIArray(t *testing.T) {
t.Error("Want value 42, got", v)
}

count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, nil)
var cursor BatchCursor
count, err = m.BatchLookup(&cursor, lookupKeys, lookupValues, nil)
if err != nil {
t.Fatalf("BatchLookup: %v", err)
}
if count != len(lookupKeys) {
t.Fatalf("BatchLookup: returned %d results, expected %d", count, len(lookupKeys))
}
if nextKey != lookupKeys[1] {
t.Fatalf("BatchLookup: expected nextKey, %d, to be the same as the lastKey returned, %d", nextKey, lookupKeys[1])
}
if !reflect.DeepEqual(keys, lookupKeys) {
t.Errorf("BatchUpdate and BatchLookup keys disagree: %v %v", keys, lookupKeys)
}
if !reflect.DeepEqual(values, lookupValues) {
t.Errorf("BatchUpdate and BatchLookup values disagree: %v %v", values, lookupValues)
}

_, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
cursor = BatchCursor{}
_, err = m.BatchLookupAndDelete(&cursor, deleteKeys, deleteValues, nil)
if !errors.Is(err, ErrNotSupported) {
t.Fatalf("BatchLookUpDelete: expected error %v, but got %v", ErrNotSupported, err)
}
Expand All @@ -172,7 +170,6 @@ func TestBatchAPIHash(t *testing.T) {
defer m.Close()

var (
nextKey uint32
keys = []uint32{0, 1}
values = []uint32{42, 4242}
lookupKeys = make([]uint32, 2)
Expand All @@ -197,7 +194,8 @@ func TestBatchAPIHash(t *testing.T) {
t.Error("Want value 42, got", v)
}

count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, nil)
var cursor BatchCursor
count, err = m.BatchLookup(&cursor, lookupKeys, lookupValues, nil)
if !errors.Is(err, ErrKeyNotExist) {
t.Fatalf("BatchLookup: expected %v got %v", ErrKeyNotExist, err)
}
Expand All @@ -213,7 +211,8 @@ func TestBatchAPIHash(t *testing.T) {
t.Errorf("BatchUpdate and BatchLookup values disagree: %v %v", values, lookupValues)
}

count, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
cursor = BatchCursor{}
count, err = m.BatchLookupAndDelete(&cursor, deleteKeys, deleteValues, nil)
if !errors.Is(err, ErrKeyNotExist) {
t.Fatalf("BatchLookupAndDelete: expected %v got %v", ErrKeyNotExist, err)
}
Expand Down Expand Up @@ -345,21 +344,21 @@ func TestBatchMapWithLock(t *testing.T) {
t.Fatalf("BatchUpdate: expected count, %d, to be %d", count, len(keys))
}

nextKey := uint32(0)
var cursor BatchCursor
lookupKeys := make([]uint32, 2)
lookupValues := make([]spinLockValue, 2)
count, err = m.BatchLookup(nil, &nextKey, lookupKeys, lookupValues, &BatchOptions{ElemFlags: uint64(LookupLock)})
count, err = m.BatchLookup(&cursor, lookupKeys, lookupValues, &BatchOptions{ElemFlags: uint64(LookupLock)})
if !errors.Is(err, ErrKeyNotExist) {
t.Fatalf("BatchLookup: %v", err)
}
if count != 2 {
t.Fatalf("BatchLookup: expected two keys, got %d", count)
}

nextKey = uint32(0)
cursor = BatchCursor{}
deleteKeys := []uint32{0, 1}
deleteValues := make([]spinLockValue, 2)
count, err = m.BatchLookupAndDelete(nil, &nextKey, deleteKeys, deleteValues, nil)
count, err = m.BatchLookupAndDelete(&cursor, deleteKeys, deleteValues, nil)
if !errors.Is(err, ErrKeyNotExist) {
t.Fatalf("BatchLookupAndDelete: %v", err)
}
Expand Down Expand Up @@ -2108,9 +2107,9 @@ func BenchmarkIterate(b *testing.B) {

b.ReportAllocs()

var cursor BatchCursor
for i := 0; i < b.N; i++ {
var next uint32
_, err := m.BatchLookup(nil, &next, k, v, nil)
_, err := m.BatchLookup(&cursor, k, v, nil)
if err != nil && !errors.Is(err, ErrKeyNotExist) {
b.Fatal(err)
}
Expand All @@ -2130,8 +2129,8 @@ func BenchmarkIterate(b *testing.B) {
}
b.StartTimer()

var next uint32
_, err := m.BatchLookupAndDelete(nil, &next, k, v, nil)
var cursor BatchCursor
_, err := m.BatchLookupAndDelete(&cursor, k, v, nil)
if err != nil && !errors.Is(err, ErrKeyNotExist) {
b.Fatal(err)
}
Expand Down

0 comments on commit d9f8904

Please sign in to comment.