diff --git a/internal/sysenc/buffer.go b/internal/sysenc/buffer.go index c6959d9cc..97ac473c8 100644 --- a/internal/sysenc/buffer.go +++ b/internal/sysenc/buffer.go @@ -32,6 +32,7 @@ func UnsafeBuffer(ptr unsafe.Pointer) Buffer { // SyscallOutput prepares a Buffer for a syscall to write into. // +// size is the length of the desired buffer in bytes. // The buffer may point at the underlying memory of dst, in which case [Unmarshal] // becomes a no-op. // diff --git a/map.go b/map.go index 87a91d133..5e71ccfc2 100644 --- a/map.go +++ b/map.go @@ -1048,15 +1048,14 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i if count != valuesValue.Len() { return 0, fmt.Errorf("keysOut and valuesOut must be the same length") } - keyBuf := make([]byte, count*int(m.keySize)) - keyPtr := sys.NewSlicePointer(keyBuf) - valueBuf := make([]byte, count*int(m.fullValueSize)) - valuePtr := sys.NewSlicePointer(valueBuf) + + keyBuf := sysenc.SyscallOutput(keysOut, count*int(m.keySize)) + valueBuf := sysenc.SyscallOutput(valuesOut, count*int(m.fullValueSize)) attr := sys.MapLookupBatchAttr{ MapFd: m.fd.Uint(), - Keys: keyPtr, - Values: valuePtr, + Keys: keyBuf.Pointer(), + Values: valueBuf.Pointer(), Count: uint32(count), InBatch: sys.NewSlicePointer(inBatch), OutBatch: sys.NewSlicePointer(cursor.opaque), @@ -1073,12 +1072,10 @@ func (m *Map) batchLookup(cmd sys.Cmd, cursor *BatchCursor, keysOut, valuesOut i return 0, sysErr } - err := sysenc.Unmarshal(keysOut, keyBuf) - if err != nil { + if err := keyBuf.Unmarshal(keysOut); err != nil { return 0, err } - err = sysenc.Unmarshal(valuesOut, valueBuf) - if err != nil { + if err := valueBuf.Unmarshal(valuesOut); err != nil { return 0, err } diff --git a/map_test.go b/map_test.go index 415953916..13e52b2a3 100644 --- a/map_test.go +++ b/map_test.go @@ -1093,9 +1093,8 @@ func TestMapIteratorAllocations(t *testing.T) { var k, v uint32 iter := arr.Iterate() - // Allocate any necessary temporary buffers. - qt.Assert(t, qt.IsTrue(iter.Next(&k, &v))) + // AllocsPerRun warms up the function for us. allocs := testing.AllocsPerRun(1, func() { if !iter.Next(&k, &v) { t.Fatal("Next failed") @@ -1105,6 +1104,35 @@ func TestMapIteratorAllocations(t *testing.T) { qt.Assert(t, qt.Equals(allocs, float64(0))) } +func TestMapBatchLookupAllocations(t *testing.T) { + testutils.SkipIfNotSupported(t, haveBatchAPI()) + + arr, err := NewMap(&MapSpec{ + Type: Array, + KeySize: 4, + ValueSize: 4, + MaxEntries: 10, + }) + if err != nil { + t.Fatal(err) + } + defer arr.Close() + + var cursor BatchCursor + tmp := make([]uint32, 2) + input := any(tmp) + + // AllocsPerRun warms up the function for us. + allocs := testing.AllocsPerRun(1, func() { + _, err := arr.BatchLookup(&cursor, input, input, nil) + if err != nil { + t.Fatal(err) + } + }) + + qt.Assert(t, qt.Equals(allocs, 0)) +} + func TestMapIterateHashKeyOneByteFull(t *testing.T) { hash, err := NewMap(&MapSpec{ Type: Hash,