From da2e07e55851a4cf321c23b55e8b9433e3b6caf5 Mon Sep 17 00:00:00 2001 From: Jesse Peterson Date: Fri, 14 Jun 2024 11:31:24 -0700 Subject: [PATCH] inventory subsystem kv storage --- subsystem/inventory/storage/diskv/diskv.go | 80 ++------------ .../inventory/storage/diskv/diskv_test.go | 4 +- subsystem/inventory/storage/inmem/inmem.go | 55 ++-------- subsystem/inventory/storage/kv/kv.go | 101 ++++++++++++++++++ subsystem/inventory/storage/storage.go | 12 +++ subsystem/inventory/storage/test/test.go | 2 +- 6 files changed, 131 insertions(+), 123 deletions(-) create mode 100644 subsystem/inventory/storage/kv/kv.go diff --git a/subsystem/inventory/storage/diskv/diskv.go b/subsystem/inventory/storage/diskv/diskv.go index 1f79a93..9e6a7ea 100644 --- a/subsystem/inventory/storage/diskv/diskv.go +++ b/subsystem/inventory/storage/diskv/diskv.go @@ -1,87 +1,27 @@ -// Package diskv implements a diskv-backed inventory subsystem storage backend. +// Package diskv implements an inventory subsystem backend using diskv. package diskv import ( - "context" - "encoding/json" - "fmt" "path/filepath" - "github.com/micromdm/nanocmd/subsystem/inventory/storage" + "github.com/micromdm/nanocmd/subsystem/inventory/storage/kv" + + "github.com/micromdm/nanolib/storage/kv/kvdiskv" "github.com/peterbourgon/diskv/v3" ) -// Diskv is an on-disk enrollment inventory data store. +// Diskv is an inventory subsystem backend which uses diskv as the key-value store. type Diskv struct { - diskv *diskv.Diskv + *kv.KV } -// New creates a new initialized inventory data store. +// New creates a new profile store at on disk at path. func New(path string) *Diskv { - flatTransform := func(s string) []string { return []string{} } return &Diskv{ - diskv: diskv.New(diskv.Options{ + KV: kv.New(kvdiskv.New(diskv.New(diskv.Options{ BasePath: filepath.Join(path, "inventory"), - Transform: flatTransform, + Transform: kvdiskv.FlatTransform, CacheSizeMax: 1024 * 1024, - }), - } -} - -// RetrieveInventory retrieves the inventory data for enrollment IDs. -func (s *Diskv) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) { - ret := make(map[string]storage.Values) - for _, id := range opt.IDs { - if !s.diskv.Has(id) { - continue - } - raw, err := s.diskv.Read(id) - if err != nil { - return ret, fmt.Errorf("reading values for %s: %w", id, err) - } - var vals storage.Values - if err = json.Unmarshal(raw, &vals); err != nil { - return ret, fmt.Errorf("unmarshal values for %s: %w", id, err) - } - ret[id] = vals - } - return ret, nil -} - -// StoreInventoryValues stores inventory data about the specified ID. -func (s *Diskv) StoreInventoryValues(ctx context.Context, id string, values storage.Values) error { - var err error - var raw []byte - var vals storage.Values - if s.diskv.Has(id) { - // this is likely race-prone as we perform a read-process-write on the same key. - if raw, err = s.diskv.Read(id); err != nil { - return fmt.Errorf("reading values: %w", err) - } - if len(raw) > 0 { - if err = json.Unmarshal(raw, &vals); err != nil { - return fmt.Errorf("unmarshal values: %w", err) - } - if vals != nil { - for k := range values { - vals[k] = values[k] - } - } - } - } - if vals == nil { - vals = values - } - if raw, err = json.Marshal(vals); err != nil { - return fmt.Errorf("marshal values: %w", err) + }))), } - if err = s.diskv.Write(id, raw); err != nil { - return fmt.Errorf("write values: %w", err) - } - return nil -} - -// DeleteInventory deletes all inventory data for an enrollment ID. -func (s *Diskv) DeleteInventory(ctx context.Context, id string) error { - return s.diskv.Erase(id) } diff --git a/subsystem/inventory/storage/diskv/diskv_test.go b/subsystem/inventory/storage/diskv/diskv_test.go index 263b480..0817aa3 100644 --- a/subsystem/inventory/storage/diskv/diskv_test.go +++ b/subsystem/inventory/storage/diskv/diskv_test.go @@ -1,7 +1,6 @@ package diskv import ( - "os" "testing" "github.com/micromdm/nanocmd/subsystem/inventory/storage" @@ -9,6 +8,5 @@ import ( ) func TestDiskv(t *testing.T) { - test.TestStorage(t, func() storage.Storage { return New("teststor") }) - os.RemoveAll("teststor") + test.TestStorage(t, func() storage.Storage { return New(t.TempDir()) }) } diff --git a/subsystem/inventory/storage/inmem/inmem.go b/subsystem/inventory/storage/inmem/inmem.go index 605682e..e8412f5 100644 --- a/subsystem/inventory/storage/inmem/inmem.go +++ b/subsystem/inventory/storage/inmem/inmem.go @@ -2,60 +2,17 @@ package inmem import ( - "context" - "sync" + "github.com/micromdm/nanocmd/subsystem/inventory/storage/kv" - "github.com/micromdm/nanocmd/subsystem/inventory/storage" + "github.com/micromdm/nanolib/storage/kv/kvmap" ) -// InMem represents the in-memory enrollment inventory data store. +// InMem is an in-memory inventory subsystem storage system backend. type InMem struct { - mu sync.RWMutex - inv map[string]storage.Values + *kv.KV } -// New creates a new initialized inventory data store. +// New creates a new inventory subsystem storage system backend. func New() *InMem { - return &InMem{inv: make(map[string]storage.Values)} -} - -// RetrieveInventory retrieves the inventory data for enrollment IDs. -func (s *InMem) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) { - s.mu.RLock() - defer s.mu.RUnlock() - if opt == nil || len(opt.IDs) <= 0 { - return nil, nil - } - ret := make(map[string]storage.Values) - for _, id := range opt.IDs { - if vals, ok := s.inv[id]; ok { - ret[id] = make(storage.Values) - for k, v := range vals { - ret[id][k] = v - } - } - } - return ret, nil -} - -// StoreInventoryValues stores inventory data about the specified ID. -func (s *InMem) StoreInventoryValues(ctx context.Context, id string, values storage.Values) error { - s.mu.Lock() - defer s.mu.Unlock() - if s.inv[id] == nil { - s.inv[id] = values - } else { - for k, v := range values { - s.inv[id][k] = v - } - } - return nil -} - -// DeleteInventory deletes all inventory data for an enrollment ID. -func (s *InMem) DeleteInventory(ctx context.Context, id string) error { - s.mu.Lock() - defer s.mu.Unlock() - delete(s.inv, id) - return nil + return &InMem{KV: kv.New(kvmap.New())} } diff --git a/subsystem/inventory/storage/kv/kv.go b/subsystem/inventory/storage/kv/kv.go new file mode 100644 index 0000000..6b56dc3 --- /dev/null +++ b/subsystem/inventory/storage/kv/kv.go @@ -0,0 +1,101 @@ +// Package kv implements an inventory subsystem storage backend using a key-value store. +package kv + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + + "github.com/micromdm/nanocmd/subsystem/inventory/storage" + + "github.com/micromdm/nanolib/storage/kv" +) + +// KV is an inventory subsystem storage backend using a key-value store. +type KV struct { + mu sync.RWMutex + b kv.KeysPrefixTraversingBucket +} + +// New creates a new inventory subsystem backend. +func New(b kv.KeysPrefixTraversingBucket) *KV { + return &KV{b: b} +} + +// RetrieveInventory queries and returns the inventory values by mapped +// by enrollment ID from the key-value store. Must provide opt and IDs. +func (s *KV) RetrieveInventory(ctx context.Context, opt *storage.SearchOptions) (map[string]storage.Values, error) { + if opt == nil || len(opt.IDs) < 1 { + return nil, storage.ErrNoIDs + } + + s.mu.RLock() + defer s.mu.RUnlock() + + r := make(map[string]storage.Values) + for _, id := range opt.IDs { + jsonValues, err := s.b.Get(ctx, id) + if errors.Is(err, kv.ErrKeyNotFound) { + continue + } else if err != nil { + return r, fmt.Errorf("getting values for %s: %w", id, err) + } + + var values storage.Values + if err = json.Unmarshal(jsonValues, &values); err != nil { + return r, fmt.Errorf("unmarshal values for %s: %w", id, err) + } + r[id] = values + } + return r, nil +} + +// StoreInventoryValues stores inventory data about the specified ID. +func (s *KV) StoreInventoryValues(ctx context.Context, id string, newValues storage.Values) error { + if id == "" { + return storage.ErrNoIDs + } + if len(newValues) == 0 { + return nil + } + + s.mu.Lock() + defer s.mu.Unlock() + + jsonValues, err := s.b.Get(ctx, id) + if err != nil && !errors.Is(err, kv.ErrKeyNotFound) { + return fmt.Errorf("get values: %w", err) + } + + var values storage.Values + if len(jsonValues) < 1 { + values = newValues + } else { + // load existing values + if err = json.Unmarshal(jsonValues, &values); err != nil { + return fmt.Errorf("unmarshal values: %w", err) + } + + // merge the new values in + for k := range newValues { + values[k] = newValues[k] + } + } + + if jsonValues, err = json.Marshal(&values); err != nil { + return fmt.Errorf("marshal values: %w", err) + } + + if err = s.b.Set(ctx, id, jsonValues); err != nil { + return fmt.Errorf("set values: %w", err) + } + + return nil +} + +// DeleteInventory deletes all inventory data for an enrollment ID. +func (s *KV) DeleteInventory(ctx context.Context, id string) error { + return s.b.Delete(ctx, id) +} diff --git a/subsystem/inventory/storage/storage.go b/subsystem/inventory/storage/storage.go index 211b1be..9327613 100644 --- a/subsystem/inventory/storage/storage.go +++ b/subsystem/inventory/storage/storage.go @@ -3,6 +3,11 @@ package storage import ( "context" + "errors" +) + +var ( + ErrNoIDs = errors.New("no ids supplied") ) // SearchOptions is a basic query for inventory of enrollment IDs. @@ -15,11 +20,18 @@ type Values map[string]interface{} type ReadStorage interface { // RetrieveInventory queries and returns the inventory values by mapped by enrollment ID. + // If no search opt nor IDs are provided an ErrNoIDs should be returned. + // If IDs are have no inventory data then they should be skipped and + // omitted from the output with no error. RetrieveInventory(ctx context.Context, opt *SearchOptions) (map[string]Values, error) } type Storage interface { ReadStorage + + // StoreInventoryValues stores inventory data about the specified ID. StoreInventoryValues(ctx context.Context, id string, values Values) error + + // DeleteInventory deletes all inventory data for an enrollment ID. DeleteInventory(ctx context.Context, id string) error } diff --git a/subsystem/inventory/storage/test/test.go b/subsystem/inventory/storage/test/test.go index 053a5da..c7cf4d3 100644 --- a/subsystem/inventory/storage/test/test.go +++ b/subsystem/inventory/storage/test/test.go @@ -57,6 +57,6 @@ func TestStorage(t *testing.T, newStorage func() storage.Storage) { _, ok = idVals[id] if ok { - t.Error("expected id to missing in id values map") + t.Error("expected id to be missing in id values map") } }