lpm_index_test.go

// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package statedb

import (
	"encoding/binary"
	"fmt"
	"iter"
	"math/bits"
	"math/rand"
	"net/netip"
	"reflect"
	"slices"
	"testing"
	"testing/quick"

	"github.com/cilium/statedb/index"
	"github.com/cilium/statedb/lpm"
	"github.com/stretchr/testify/require"
)

type lpmTestObject struct {
	ID            uint16
	Prefix        netip.Prefix
	Port          uint16
	PortPrefixLen lpm.PrefixLen
}

// TableHeader implements TableWritable.
func (l lpmTestObject) TableHeader() []string {
	return []string{
		"ID",
		"Prefix",
	}
}

// TableRow implements TableWritable.
func (l lpmTestObject) TableRow() []string {
	return []string{
		fmt.Sprintf("%d", l.ID),
		l.Prefix.String(),
	}
}

var _ TableWritable = lpmTestObject{}

var (
	lpmIDIndex = Index[lpmTestObject, uint16]{
		Name: "id",
		FromObject: func(obj lpmTestObject) index.KeySet {
			return index.NewKeySet(index.Uint16(obj.ID))
		},
		FromKey:    index.Uint16,
		FromString: index.Uint16String,
		Unique:     true,
	}

	lpmPrefixIndex = NetIPPrefixIndex[lpmTestObject]{
		Name: "prefix",
		FromObject: func(obj lpmTestObject) iter.Seq[netip.Prefix] {
			return Just(obj.Prefix)
		},
		Unique: true,
	}

	lpmPortIndex = LPMIndex[lpmTestObject]{
		Name:   "port",
		Unique: true,
		FromObject: func(obj lpmTestObject) iter.Seq2[[]byte, lpm.PrefixLen] {
			return Just2(
				binary.BigEndian.AppendUint16(nil, obj.Port),
				obj.PortPrefixLen)
		},
	}

	lpmPortNonUniqueIndex = LPMIndex[lpmTestObject]{
		Name:   "port-non-unique",
		Unique: false,
		FromObject: func(obj lpmTestObject) iter.Seq2[[]byte, lpm.PrefixLen] {
			return Just2(
				binary.BigEndian.AppendUint16(nil, obj.Port),
				obj.PortPrefixLen)
		},
	}
)

func newLPMTestTable(db *DB) RWTable[lpmTestObject] {
	t, err := NewTable(
		db,
		"lpm-test",
		lpmIDIndex,
		lpmPrefixIndex,
		lpmPortIndex,
	)
	if err != nil {
		panic(err)
	}
	return t
}

func newLPMNonUniqueTestTable(db *DB) RWTable[lpmTestObject] {
	t, err := NewTable(
		db,
		"lpm-test-non-unique",
		lpmIDIndex,
		lpmPrefixIndex,
		lpmPortNonUniqueIndex,
	)
	if err != nil {
		panic(err)
	}
	return t
}

func TestLPMIndex(t *testing.T) {
	db := New()
	tbl := newLPMTestTable(db)

	wtxn := db.WriteTxn(tbl)
	tbl.Insert(wtxn, lpmTestObject{
		ID:            0,
		Prefix:        netip.MustParsePrefix("1.0.0.0/8"),
		PortPrefixLen: 16,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            1,
		Prefix:        netip.MustParsePrefix("10.0.0.0/8"),
		Port:          0xff00,
		PortPrefixLen: 8,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            2,
		Prefix:        netip.MustParsePrefix("10.0.0.0/24"),
		PortPrefixLen: 16,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            3,
		Prefix:        netip.MustParsePrefix("10.0.0.1/32"),
		PortPrefixLen: 16,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            4,
		Prefix:        netip.MustParsePrefix("192.168.0.0/24"),
		PortPrefixLen: 16,
	})
	txn := wtxn.Commit()

	obj, _, found := tbl.Get(txn, lpmPortIndex.Query([]byte{0xff, 0x11}, 16))
	require.True(t, found)
	require.EqualValues(t, 1, obj.ID)

	obj, _, found = tbl.Get(txn, lpmPrefixIndex.Query(netip.MustParseAddr("10.0.1.1")))
	require.True(t, found)
	require.EqualValues(t, 1, obj.ID)

	objs := Collect(tbl.Prefix(txn, lpmPrefixIndex.QueryPrefix(netip.MustParsePrefix("10.0.0.0/8"))))
	require.Len(t, objs, 3)
	require.EqualValues(t, objs[0].ID, 1)
	require.EqualValues(t, objs[1].ID, 2)
	require.EqualValues(t, objs[2].ID, 3)

	objs = Collect(tbl.LowerBound(txn, lpmPrefixIndex.QueryPrefix(netip.MustParsePrefix("10.0.0.0/8"))))
	require.Len(t, objs, 4)
	require.EqualValues(t, objs[0].ID, 1)
	require.EqualValues(t, objs[1].ID, 2)
	require.EqualValues(t, objs[2].ID, 3)
	require.EqualValues(t, objs[3].ID, 4)

	// Since the LPM indexes are unique List() is the same as Get()
	objs = Collect(tbl.List(txn, lpmPrefixIndex.QueryPrefix(netip.MustParsePrefix("10.0.0.0/8"))))
	require.Len(t, objs, 1)
	require.EqualValues(t, objs[0].ID, 1)

	wtxn = db.WriteTxn(tbl)
	old, found, _ := tbl.Delete(wtxn, lpmTestObject{ID: 4})
	require.True(t, found)
	require.EqualValues(t, 4, old.ID)
	old, found, _ = tbl.Delete(wtxn, lpmTestObject{ID: 3})
	require.True(t, found)
	require.EqualValues(t, 3, old.ID)
	old, found, _ = tbl.Delete(wtxn, lpmTestObject{ID: 2})
	require.True(t, found)
	require.EqualValues(t, 2, old.ID)
	old, found, _ = tbl.Delete(wtxn, lpmTestObject{ID: 1})
	require.True(t, found)
	require.EqualValues(t, 1, old.ID)
	old, found, _ = tbl.Delete(wtxn, lpmTestObject{ID: 0})
	require.True(t, found)
	require.EqualValues(t, 0, old.ID)
	txn = wtxn.Commit()

	wtxn = db.WriteTxn(tbl)
	tbl.Insert(wtxn, lpmTestObject{
		ID:     4,
		Prefix: netip.MustParsePrefix("192.168.0.0/24"),
	})
	txn = wtxn.Commit()

}

func TestLPMIndexNonUnique(t *testing.T) {
	db := New()
	tbl := newLPMNonUniqueTestTable(db)

	wtxn := db.WriteTxn(tbl)
	key := binary.BigEndian.AppendUint16(nil, 0x1234)
	prefix := netip.MustParsePrefix("10.0.0.0/8")
	tbl.Insert(wtxn, lpmTestObject{
		ID:            10,
		Prefix:        prefix,
		Port:          0x1234,
		PortPrefixLen: 16,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            20,
		Prefix:        prefix.Masked(),
		Port:          0x1234,
		PortPrefixLen: 16,
	})
	tbl.Insert(wtxn, lpmTestObject{
		ID:            30,
		Prefix:        netip.MustParsePrefix("10.1.0.0/16"),
		Port:          0x1200,
		PortPrefixLen: 8,
	})
	txn := wtxn.Commit()

	// Get returns the first matching object.
	obj, _, found := tbl.Get(txn, lpmPortNonUniqueIndex.Query(key, 16))
	require.True(t, found)
	require.EqualValues(t, 10, obj.ID)

	// List returns all matches for the key.
	objs := Collect(tbl.List(txn, lpmPortNonUniqueIndex.Query(key, 16)))
	require.Len(t, objs, 2)
	require.EqualValues(t, []uint16{10, 20}, []uint16{objs[0].ID, objs[1].ID})

	// Prefix iteration returns all objects matching the broader prefix.
	objs = Collect(tbl.Prefix(txn, lpmPortNonUniqueIndex.Query(key, 8)))
	require.Len(t, objs, 3)
	require.ElementsMatch(t, []uint16{10, 20, 30}, []uint16{objs[0].ID, objs[1].ID, objs[2].ID})

	// Deleting one object keeps the other accessible.
	wtxn = db.WriteTxn(tbl)
	old, found, _ := tbl.Delete(wtxn, lpmTestObject{ID: 10})
	require.True(t, found)
	require.EqualValues(t, 10, old.ID)
	txn = wtxn.Commit()

	objs = Collect(tbl.List(txn, lpmPortNonUniqueIndex.Query(key, 16)))
	require.Len(t, objs, 1)
	require.EqualValues(t, 20, objs[0].ID)
}

func TestLPMIndexIteratorNonUnique(t *testing.T) {
	idx := lpmIndex{
		lpm:          lpm.New[lpmEntry](),
		unique:       false,
		objectToKeys: func(object) index.KeySet { return index.KeySet{} },
		watch:        make(chan struct{}),
	}
	txn := lpmIndexTxn{
		index: idx,
		tx:    idx.lpm.Txn(),
		size:  0,
	}

	obj := func(id uint16, port uint16, prefixLen lpm.PrefixLen, prefix netip.Prefix) object {
		return object{
			data: lpmTestObject{
				ID:            id,
				Prefix:        prefix,
				Port:          port,
				PortPrefixLen: prefixLen,
			},
			revision: uint64(id),
		}
	}

	key16 := lpm.EncodeLPMKey(binary.BigEndian.AppendUint16(nil, 0x1234), 16)
	key8 := lpm.EncodeLPMKey(binary.BigEndian.AppendUint16(nil, 0x1200), 8)

	txn.insertKey(index.Uint16(10), key16, obj(10, 0x1234, 16, netip.MustParsePrefix("10.0.0.0/8")))
	txn.insertKey(index.Uint16(20), key16, obj(20, 0x1234, 16, netip.MustParsePrefix("10.0.0.0/8")))
	txn.insertKey(index.Uint16(30), key8, obj(30, 0x1200, 8, netip.MustParsePrefix("10.1.0.0/16")))

	iter, _ := txn.prefix(lpm.EncodeLPMKey(binary.BigEndian.AppendUint16(nil, 0x1200), 8))
	var ids []uint16
	iter.All(func(_ []byte, obj object) bool {
		ids = append(ids, obj.data.(lpmTestObject).ID)
		return true
	})

	require.Equal(t, []uint16{30, 10, 20}, ids)
}

func TestQuickLPMIndex(t *testing.T) {
	db := New()
	tbl := newLPMTestTable(db)
	values := map[uint16]netip.Prefix{}

	check := func(id uint16, prefix netip.Prefix, remove bool) bool {
		oldPrefix, oldExists := values[id]
		if oldExists {
			obj, _, found := tbl.Get(db.ReadTxn(), lpmIDIndex.Query(id))
			if !found {
				t.Logf("inserted object not found")
				return false
			}
			if obj.Prefix != oldPrefix {
				t.Logf("object has mismatching prefix")
				return false
			}
		}

		txn := db.WriteTxn(tbl)
		defer txn.Abort()
		if remove {
			oldObj, found, _ := tbl.Delete(txn, lpmTestObject{
				ID: id,
			})
			if oldExists != found {
				t.Logf("Delete did not return the old object")
				return false
			}
			if oldPrefix != oldObj.Prefix {
				t.Logf("prefix did not match the returned one")
			}
			txn.Commit()
			delete(values, id)
			return true
		}

		tbl.Insert(txn, lpmTestObject{
			ID:     id,
			Prefix: prefix,
		})
		values[id] = prefix

		check := func(rtxn ReadTxn) bool {
			obj, _, found := tbl.Get(rtxn, lpmIDIndex.Query(id))
			if !found {
				t.Logf("inserted object not found")
				return false
			}
			if obj.ID != id || obj.Prefix != prefix {
				t.Logf("object returned with wrong data")
				return false
			}

			obj, _, found = tbl.Get(rtxn, lpmPrefixIndex.QueryPrefix(prefix))
			if !found {
				t.Logf("object not found by prefix")
				return false
			}
			if obj.ID != id || obj.Prefix != prefix {
				t.Logf("object returned with wrong data")
				return false
			}

			for obj := range tbl.Prefix(rtxn, lpmPrefixIndex.QueryPrefix(prefix)) {
				if !prefix.Overlaps(obj.Prefix) {
					t.Logf("Prefix() returned object that was not under queried prefix")
					return false
				}
				if obj.Prefix.Bits() < prefix.Bits() {
					t.Logf("Prefix() returned object with wider prefix: %s is wider than %s", obj.Prefix, prefix)
					return false
				}
			}

			for obj := range tbl.LowerBound(rtxn, lpmPrefixIndex.Query(prefix.Addr())) {
				if obj.Prefix.Addr().Compare(prefix.Addr()) < 0 {
					t.Logf("LowerBound() returned object whose address was less than query address")
				}
			}

			return true
		}
		// Try queries against the WriteTxn first
		if !check(txn) {
			t.Logf("query failed against WriteTxn")
			return false
		}
		// And once more after commmitting with a ReadTxn
		return check(txn.Commit())
	}
	require.NoError(t, quick.Check(check, &quick.Config{
		MaxCount: 5000,

		// Make 1-8 byte long strings as input data. Keep the strings shorter
		// than the default quick value generation to hit the more interesting cases
		// often.
		Values: func(args []reflect.Value, rand *rand.Rand) {
			if len(args) != 3 {
				panic("unexpected args count")
			}
			bits := rand.Intn(32)
			var addrBytes [4]byte
			rand.Read(addrBytes[:])
			addr := netip.AddrFrom4(addrBytes)
			args[0] = reflect.ValueOf(uint16(rand.Intn(65535)))
			args[1] = reflect.ValueOf(netip.PrefixFrom(addr, bits).Masked())
			args[2] = reflect.ValueOf(rand.Intn(2) == 1)
		},
	}))
}

func TestQuickLPMIndexNonUnique(t *testing.T) {
	db := New()
	tbl := newLPMNonUniqueTestTable(db)
	values := map[uint16]lpmTestObject{}

	toIDs := func(objs []lpmTestObject) []uint16 {
		if len(objs) == 0 {
			return nil
		}
		ids := make([]uint16, len(objs))
		for i, obj := range objs {
			ids[i] = obj.ID
		}
		return ids
	}

	matchPrefix := func(port, candidate uint16, length lpm.PrefixLen) (bool, lpm.PrefixLen) {
		if length == 0 {
			return true, 0
		}
		mask := uint16(0xffff) << (16 - length)
		x := port&mask ^ candidate&mask
		return (port & mask) == (candidate & mask), min(length, lpm.PrefixLen(bits.LeadingZeros16(x)))
	}

	longestMatch := func(port uint16) (ids []uint16) {
		matchLen := lpm.PrefixLen(0)
		for id, obj := range values {
			match, numBits := matchPrefix(port, obj.Port, obj.PortPrefixLen)
			if match && numBits >= matchLen {
				if numBits > matchLen {
					matchLen = numBits
					ids = nil
				}
				ids = append(ids, id)
			}
		}
		slices.Sort(ids)
		return
	}

	validate := func(rtxn ReadTxn, port uint16) bool {
		if len(values) != tbl.NumObjects(rtxn) {
			t.Logf("wrong object count, expected %d got %d", len(values), tbl.NumObjects(rtxn))
			return false
		}

		key := binary.BigEndian.AppendUint16(nil, port)
		listObjs := Collect(tbl.List(rtxn, lpmPortNonUniqueIndex.Query(key, 16)))
		expectedIDs := longestMatch(port)
		listIDs := toIDs(listObjs)
		if !slices.Equal(expectedIDs, listIDs) {
			t.Logf("List(%d/16) returned %v, expected %v", port, listIDs, expectedIDs)
			return false
		}

		obj, _, found := tbl.Get(rtxn, lpmPortNonUniqueIndex.Query(key, 16))
		if len(expectedIDs) > 0 && !found {
			t.Logf("expected result but Get returned no object")
			return false
		}
		if found && !slices.Contains(expectedIDs, obj.ID) {
			t.Logf("Get(%d/16) returned object %d that was not part of expected IDs %v", port, obj.ID, expectedIDs)
			return false
		}
		return true
	}

	check := func(id uint16, prefix netip.Prefix, port uint16, portPrefix uint8, remove bool) bool {
		prefix = prefix.Masked()
		portPrefixLen := lpm.PrefixLen(portPrefix % 17)

		current, exists := values[id]
		if exists {
			obj, _, found := tbl.Get(db.ReadTxn(), lpmIDIndex.Query(id))
			if !found {
				t.Logf("existing object not found by id")
				return false
			}
			if obj.Prefix != current.Prefix || obj.Port != current.Port || obj.PortPrefixLen != current.PortPrefixLen {
				t.Logf("existing object does not match stored data")
				return false
			}
		}

		txn := db.WriteTxn(tbl)
		defer txn.Abort()
		if remove {
			oldObj, found, _ := tbl.Delete(txn, lpmTestObject{ID: id})
			if found != exists {
				t.Logf("delete found=%v, expected %v", found, exists)
				return false
			}
			if found {
				if oldObj.ID != id {
					t.Logf("delete returned wrong id %d vs %d", oldObj.ID, id)
					return false
				}
				if oldObj.Prefix != current.Prefix || oldObj.Port != current.Port || oldObj.PortPrefixLen != current.PortPrefixLen {
					t.Logf("delete returned object %v, expected %v", oldObj, current)
					return false
				}
				delete(values, id)
			}
			txn.Commit()
			return validate(db.ReadTxn(), port)
		}

		obj := lpmTestObject{
			ID:            id,
			Prefix:        prefix,
			Port:          port,
			PortPrefixLen: portPrefixLen,
		}
		oldObj, hadOld, err := tbl.Insert(txn, obj)
		if err != nil {
			t.Logf("insert error: %s", err)
			return false
		}
		if hadOld != exists {
			t.Logf("Insert returned hadOld=%v, expected %v", hadOld, exists)
			return false
		}
		if hadOld {
			if oldObj.ID != id {
				t.Logf("Insert returned old object with wrong id %d vs %d", oldObj.ID, id)
				return false
			}
			if oldObj.Prefix != current.Prefix || oldObj.Port != current.Port || oldObj.PortPrefixLen != current.PortPrefixLen {
				t.Logf("Insert returned wrong old object %+v vs %+v", oldObj, current)
				return false
			}
		}
		values[id] = obj

		if !validate(txn, port) {
			t.Logf("validation failed against WriteTxn")
			return false
		}
		return validate(txn.Commit(), port)
	}

	require.NoError(t, quick.Check(check, &quick.Config{
		MaxCount: 5000,
		Values: func(args []reflect.Value, rand *rand.Rand) {
			if len(args) != 5 {
				panic("unexpected args count")
			}
			bits := rand.Intn(32)
			var addrBytes [4]byte
			rand.Read(addrBytes[:])
			args[0] = reflect.ValueOf(uint16(rand.Intn(65536)))
			args[1] = reflect.ValueOf(netip.PrefixFrom(netip.AddrFrom4(addrBytes), bits).Masked())
			args[2] = reflect.ValueOf(uint16(rand.Intn(65536)))
			args[3] = reflect.ValueOf(uint8(rand.Intn(17)))
			args[4] = reflect.ValueOf(rand.Intn(3) == 1)
		},
	}))
}

func BenchmarkDB_WriteTxn_100_LPMIndex(b *testing.B) {
	db := New()
	tbl := newLPMTestTable(db)
	batchSize := 100
	n := b.N

	var addrs []netip.Prefix
	for i := range n {
		addr, _ := netip.AddrFromSlice(binary.BigEndian.AppendUint32(nil, uint32(i)))
		addrs = append(addrs, netip.PrefixFrom(addr, 32))
	}
	b.ResetTimer()

	for n > 0 {
		txn := db.WriteTxn(tbl)
		toWrite := min(n, batchSize)
		for i := range toWrite {
			_, _, err := tbl.Insert(txn,
				lpmTestObject{
					ID:     uint16(i),
					Prefix: addrs[i],
				})
			if err != nil {
				b.Fatalf("Insert error: %s", err)
			}
		}
		txn.Commit()
		n -= toWrite
	}

	b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "objects/sec")
}

func BenchmarkDB_WriteTxn_1_LPMIndex(b *testing.B) {
	db := New()
	tbl := newLPMTestTable(db)

	var addrs []netip.Prefix
	for i := range b.N {
		addr, _ := netip.AddrFromSlice(binary.LittleEndian.AppendUint32(nil, uint32(i)))
		addrs = append(addrs, netip.PrefixFrom(addr, 32))
	}
	b.ResetTimer()

	for i := range b.N {
		txn := db.WriteTxn(tbl)
		_, _, err := tbl.Insert(txn,
			lpmTestObject{
				ID:     uint16(i),
				Prefix: addrs[i],
			})
		if err != nil {
			b.Fatalf("Insert error: %s", err)
		}
		txn.Commit()
	}

	b.ReportMetric(float64(b.N)/b.Elapsed().Seconds(), "objects/sec")
}

func BenchmarkDB_LPMIndex_Get(b *testing.B) {
	db := New()
	tbl := newLPMTestTable(db)
	wtxn := db.WriteTxn(tbl)
	queries := []Query[lpmTestObject]{}
	for i := range 10000 {
		addr, _ := netip.AddrFromSlice(binary.BigEndian.AppendUint32(nil, uint32(i)))
		queries = append(queries, lpmPrefixIndex.Query(addr))
		tbl.Insert(wtxn,
			lpmTestObject{
				ID:     uint16(i),
				Prefix: netip.PrefixFrom(addr, 32),
			})
	}
	txn := wtxn.Commit()
	b.ResetTimer()

	for b.Loop() {
		for _, q := range queries {
			_, _, found := tbl.Get(txn, q)
			if !found {
				b.Fatalf("not found")
			}
		}
	}
	b.ReportMetric(float64(b.N*len(queries))/b.Elapsed().Seconds(), "objects/sec")
}