412 lines
11 KiB
Go
412 lines
11 KiB
Go
package testing
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"reflect"
|
|
"sort"
|
|
"testing"
|
|
|
|
"github.com/influxdata/influxdb/v2/kv"
|
|
)
|
|
|
|
const (
|
|
someResourceBucket = "aresource"
|
|
)
|
|
|
|
var (
|
|
mapping = kv.NewIndexMapping([]byte(someResourceBucket), []byte("aresourcebyowneridv1"), func(body []byte) ([]byte, error) {
|
|
var resource someResource
|
|
if err := json.Unmarshal(body, &resource); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return []byte(resource.OwnerID), nil
|
|
})
|
|
)
|
|
|
|
type someResource struct {
|
|
ID string
|
|
OwnerID string
|
|
}
|
|
|
|
type someResourceStore struct {
|
|
store kv.Store
|
|
ownerIDIndex *kv.Index
|
|
}
|
|
|
|
func newSomeResourceStore(ctx context.Context, store kv.Store) *someResourceStore {
|
|
return &someResourceStore{
|
|
store: store,
|
|
ownerIDIndex: kv.NewIndex(mapping),
|
|
}
|
|
}
|
|
|
|
func (s *someResourceStore) FindByOwner(ctx context.Context, ownerID string) (resources []someResource, err error) {
|
|
err = s.store.View(ctx, func(tx kv.Tx) error {
|
|
return s.ownerIDIndex.Walk(ctx, tx, []byte(ownerID), func(k, v []byte) error {
|
|
var resource someResource
|
|
if err := json.Unmarshal(v, &resource); err != nil {
|
|
return err
|
|
}
|
|
|
|
resources = append(resources, resource)
|
|
return nil
|
|
})
|
|
})
|
|
return
|
|
}
|
|
|
|
func (s *someResourceStore) Create(ctx context.Context, resource someResource, index bool) error {
|
|
return s.store.Update(ctx, func(tx kv.Tx) error {
|
|
bkt, err := tx.Bucket(mapping.SourceBucket())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if index {
|
|
if err := s.ownerIDIndex.Insert(tx, []byte(resource.OwnerID), []byte(resource.ID)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
data, err := json.Marshal(resource)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return bkt.Put([]byte(resource.ID), data)
|
|
})
|
|
}
|
|
|
|
func newResource(id, owner string) someResource {
|
|
return someResource{ID: id, OwnerID: owner}
|
|
}
|
|
|
|
func newNResources(n int) (resources []someResource) {
|
|
return newNResourcesWithUserCount(n, 5)
|
|
}
|
|
|
|
func newNResourcesWithUserCount(n, userCount int) (resources []someResource) {
|
|
for i := 0; i < n; i++ {
|
|
var (
|
|
id = fmt.Sprintf("resource %d", i)
|
|
owner = fmt.Sprintf("owner %d", i%userCount)
|
|
)
|
|
resources = append(resources, newResource(id, owner))
|
|
}
|
|
return
|
|
}
|
|
|
|
func TestIndex(t *testing.T, store kv.Store) {
|
|
t.Run("Test_PopulateAndVerify", func(t *testing.T) {
|
|
testPopulateAndVerify(t, store)
|
|
})
|
|
|
|
t.Run("Test_Walk", func(t *testing.T) {
|
|
testWalk(t, store)
|
|
})
|
|
}
|
|
|
|
func testPopulateAndVerify(t *testing.T, store kv.Store) {
|
|
var (
|
|
ctx = context.TODO()
|
|
resources = newNResources(20)
|
|
resourceStore = newSomeResourceStore(ctx, store)
|
|
)
|
|
|
|
// insert 20 resources, but only index the first half
|
|
for i, resource := range resources {
|
|
if err := resourceStore.Create(ctx, resource, i < len(resources)/2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
// check that the index is populated with only 10 items
|
|
var count int
|
|
store.View(ctx, func(tx kv.Tx) error {
|
|
kvs, err := allKVs(tx, mapping.IndexBucket())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
count = len(kvs)
|
|
|
|
return nil
|
|
})
|
|
|
|
if count > 10 {
|
|
t.Errorf("expected index to be empty, found %d items", count)
|
|
}
|
|
|
|
// ensure verify identifies the 10 missing items from the index
|
|
diff, err := resourceStore.ownerIDIndex.Verify(ctx, store)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
expected := kv.IndexDiff{
|
|
PresentInIndex: map[string]map[string]struct{}{
|
|
"owner 0": {"resource 0": {}, "resource 5": {}},
|
|
"owner 1": {"resource 1": {}, "resource 6": {}},
|
|
"owner 2": {"resource 2": {}, "resource 7": {}},
|
|
"owner 3": {"resource 3": {}, "resource 8": {}},
|
|
"owner 4": {"resource 4": {}, "resource 9": {}},
|
|
},
|
|
MissingFromIndex: map[string]map[string]struct{}{
|
|
"owner 0": {"resource 10": {}, "resource 15": {}},
|
|
"owner 1": {"resource 11": {}, "resource 16": {}},
|
|
"owner 2": {"resource 12": {}, "resource 17": {}},
|
|
"owner 3": {"resource 13": {}, "resource 18": {}},
|
|
"owner 4": {"resource 14": {}, "resource 19": {}},
|
|
},
|
|
}
|
|
|
|
if !reflect.DeepEqual(expected, diff) {
|
|
t.Errorf("expected %#v, found %#v", expected, diff)
|
|
}
|
|
|
|
corrupt := diff.Corrupt()
|
|
sort.Strings(corrupt)
|
|
|
|
if expected := []string{
|
|
"owner 0",
|
|
"owner 1",
|
|
"owner 2",
|
|
"owner 3",
|
|
"owner 4",
|
|
}; !reflect.DeepEqual(expected, corrupt) {
|
|
t.Errorf("expected %#v, found %#v\n", expected, corrupt)
|
|
}
|
|
|
|
// populate the missing indexes
|
|
count, err = resourceStore.ownerIDIndex.Populate(ctx, store)
|
|
if err != nil {
|
|
t.Errorf("unexpected err %v", err)
|
|
}
|
|
|
|
// ensure only 10 items were reported as being indexed
|
|
if count != 10 {
|
|
t.Errorf("expected to index 20 items, instead indexed %d items", count)
|
|
}
|
|
|
|
// check the contents of the index
|
|
var allKvs [][2][]byte
|
|
store.View(ctx, func(tx kv.Tx) (err error) {
|
|
allKvs, err = allKVs(tx, mapping.IndexBucket())
|
|
return
|
|
})
|
|
|
|
if expected := [][2][]byte{
|
|
{[]byte("owner 0/resource 0"), []byte("resource 0")},
|
|
{[]byte("owner 0/resource 10"), []byte("resource 10")},
|
|
{[]byte("owner 0/resource 15"), []byte("resource 15")},
|
|
{[]byte("owner 0/resource 5"), []byte("resource 5")},
|
|
{[]byte("owner 1/resource 1"), []byte("resource 1")},
|
|
{[]byte("owner 1/resource 11"), []byte("resource 11")},
|
|
{[]byte("owner 1/resource 16"), []byte("resource 16")},
|
|
{[]byte("owner 1/resource 6"), []byte("resource 6")},
|
|
{[]byte("owner 2/resource 12"), []byte("resource 12")},
|
|
{[]byte("owner 2/resource 17"), []byte("resource 17")},
|
|
{[]byte("owner 2/resource 2"), []byte("resource 2")},
|
|
{[]byte("owner 2/resource 7"), []byte("resource 7")},
|
|
{[]byte("owner 3/resource 13"), []byte("resource 13")},
|
|
{[]byte("owner 3/resource 18"), []byte("resource 18")},
|
|
{[]byte("owner 3/resource 3"), []byte("resource 3")},
|
|
{[]byte("owner 3/resource 8"), []byte("resource 8")},
|
|
{[]byte("owner 4/resource 14"), []byte("resource 14")},
|
|
{[]byte("owner 4/resource 19"), []byte("resource 19")},
|
|
{[]byte("owner 4/resource 4"), []byte("resource 4")},
|
|
{[]byte("owner 4/resource 9"), []byte("resource 9")},
|
|
}; !reflect.DeepEqual(allKvs, expected) {
|
|
t.Errorf("expected %#v, found %#v", expected, allKvs)
|
|
}
|
|
|
|
// remove the last 10 items from the source, but leave them in the index
|
|
store.Update(ctx, func(tx kv.Tx) error {
|
|
bkt, err := tx.Bucket(mapping.SourceBucket())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
for _, resource := range resources[10:] {
|
|
bkt.Delete([]byte(resource.ID))
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
// ensure verify identifies the last 10 items as missing from the source
|
|
diff, err = resourceStore.ownerIDIndex.Verify(ctx, store)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
expected = kv.IndexDiff{
|
|
PresentInIndex: map[string]map[string]struct{}{
|
|
"owner 0": {"resource 0": {}, "resource 5": {}, "resource 10": {}, "resource 15": {}},
|
|
"owner 1": {"resource 1": {}, "resource 6": {}, "resource 11": {}, "resource 16": {}},
|
|
"owner 2": {"resource 2": {}, "resource 7": {}, "resource 12": {}, "resource 17": {}},
|
|
"owner 3": {"resource 3": {}, "resource 8": {}, "resource 13": {}, "resource 18": {}},
|
|
"owner 4": {"resource 4": {}, "resource 9": {}, "resource 14": {}, "resource 19": {}},
|
|
},
|
|
MissingFromSource: map[string]map[string]struct{}{
|
|
"owner 0": {"resource 10": {}, "resource 15": {}},
|
|
"owner 1": {"resource 11": {}, "resource 16": {}},
|
|
"owner 2": {"resource 12": {}, "resource 17": {}},
|
|
"owner 3": {"resource 13": {}, "resource 18": {}},
|
|
"owner 4": {"resource 14": {}, "resource 19": {}},
|
|
},
|
|
}
|
|
if !reflect.DeepEqual(expected, diff) {
|
|
t.Errorf("expected %#v, found %#v", expected, diff)
|
|
}
|
|
}
|
|
|
|
func testWalk(t *testing.T, store kv.Store) {
|
|
var (
|
|
ctx = context.TODO()
|
|
resources = newNResources(20)
|
|
// configure resource store with read disabled
|
|
resourceStore = newSomeResourceStore(ctx, store)
|
|
|
|
cases = []struct {
|
|
owner string
|
|
resources []someResource
|
|
}{
|
|
{
|
|
owner: "owner 0",
|
|
resources: []someResource{
|
|
newResource("resource 0", "owner 0"),
|
|
newResource("resource 10", "owner 0"),
|
|
newResource("resource 15", "owner 0"),
|
|
newResource("resource 5", "owner 0"),
|
|
},
|
|
},
|
|
{
|
|
owner: "owner 1",
|
|
resources: []someResource{
|
|
newResource("resource 1", "owner 1"),
|
|
newResource("resource 11", "owner 1"),
|
|
newResource("resource 16", "owner 1"),
|
|
newResource("resource 6", "owner 1"),
|
|
},
|
|
},
|
|
{
|
|
owner: "owner 2",
|
|
resources: []someResource{
|
|
newResource("resource 12", "owner 2"),
|
|
newResource("resource 17", "owner 2"),
|
|
newResource("resource 2", "owner 2"),
|
|
newResource("resource 7", "owner 2"),
|
|
},
|
|
},
|
|
{
|
|
owner: "owner 3",
|
|
resources: []someResource{
|
|
newResource("resource 13", "owner 3"),
|
|
newResource("resource 18", "owner 3"),
|
|
newResource("resource 3", "owner 3"),
|
|
newResource("resource 8", "owner 3"),
|
|
},
|
|
},
|
|
{
|
|
owner: "owner 4",
|
|
resources: []someResource{
|
|
newResource("resource 14", "owner 4"),
|
|
newResource("resource 19", "owner 4"),
|
|
newResource("resource 4", "owner 4"),
|
|
newResource("resource 9", "owner 4"),
|
|
},
|
|
},
|
|
}
|
|
)
|
|
|
|
// insert all 20 resources with indexing enabled
|
|
for _, resource := range resources {
|
|
if err := resourceStore.Create(ctx, resource, true); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
for _, testCase := range cases {
|
|
found, err := resourceStore.FindByOwner(ctx, testCase.owner)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// expect resources to be empty while read path disabled disabled
|
|
if len(found) > 0 {
|
|
t.Fatalf("expected %#v to be empty", found)
|
|
}
|
|
}
|
|
|
|
// configure index read path enabled
|
|
kv.WithIndexReadPathEnabled(resourceStore.ownerIDIndex)
|
|
|
|
for _, testCase := range cases {
|
|
found, err := resourceStore.FindByOwner(ctx, testCase.owner)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if !reflect.DeepEqual(found, testCase.resources) {
|
|
t.Errorf("expected %#v, found %#v", testCase.resources, found)
|
|
}
|
|
}
|
|
}
|
|
|
|
func allKVs(tx kv.Tx, bucket []byte) (kvs [][2][]byte, err error) {
|
|
idx, err := tx.Bucket(mapping.IndexBucket())
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
cursor, err := idx.ForwardCursor(nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
defer func() {
|
|
if cerr := cursor.Close(); cerr != nil && err == nil {
|
|
err = cerr
|
|
}
|
|
}()
|
|
|
|
for k, v := cursor.Next(); k != nil; k, v = cursor.Next() {
|
|
kvs = append(kvs, [2][]byte{k, v})
|
|
}
|
|
|
|
return kvs, cursor.Err()
|
|
}
|
|
|
|
func BenchmarkIndexWalk(b *testing.B, store kv.Store, resourceCount, fetchCount int) {
|
|
var (
|
|
ctx = context.TODO()
|
|
resourceStore = newSomeResourceStore(ctx, store)
|
|
userCount = resourceCount / fetchCount
|
|
resources = newNResourcesWithUserCount(resourceCount, userCount)
|
|
)
|
|
|
|
kv.WithIndexReadPathEnabled(resourceStore.ownerIDIndex)
|
|
|
|
for _, resource := range resources {
|
|
resourceStore.Create(ctx, resource, true)
|
|
}
|
|
|
|
b.ResetTimer()
|
|
|
|
for i := 0; i < b.N; i++ {
|
|
store.View(ctx, func(tx kv.Tx) error {
|
|
return resourceStore.ownerIDIndex.Walk(ctx, tx, []byte(fmt.Sprintf("owner %d", i%userCount)), func(k, v []byte) error {
|
|
if k == nil || v == nil {
|
|
b.Fatal("entries must not be nil")
|
|
}
|
|
|
|
return nil
|
|
})
|
|
})
|
|
}
|
|
}
|