diff --git a/ee/control/consumers/keyvalueconsumer/key_value_consumer.go b/ee/control/consumers/keyvalueconsumer/key_value_consumer.go index bd752b06e..2d2a2381f 100644 --- a/ee/control/consumers/keyvalueconsumer/key_value_consumer.go +++ b/ee/control/consumers/keyvalueconsumer/key_value_consumer.go @@ -32,7 +32,7 @@ func (c *KeyValueConsumer) Update(data io.Reader) error { } // Turn the map into a slice of key, value, ... and send it to the thing storing this data - _, err := c.updater.Update(mapToSlice(kvPairs)...) + _, err := c.updater.Update(kvPairs) return err } diff --git a/pkg/agent/flags/flag_controller.go b/pkg/agent/flags/flag_controller.go index 9913598a7..3f973d04e 100644 --- a/pkg/agent/flags/flag_controller.go +++ b/pkg/agent/flags/flag_controller.go @@ -65,14 +65,16 @@ func (fc *FlagController) set(key keys.FlagKey, value []byte) error { // Update bulk replaces agent flags and stores them. // Observers will be notified of changed flags and deleted flags. -func (fc *FlagController) Update(pairs ...string) ([]string, error) { +func (fc *FlagController) Update(kvPairs map[string]string) ([]string, error) { // Attempt to bulk replace the store with the key-values - deletedKeys, err := fc.agentFlagsStore.Update(pairs...) + deletedKeys, err := fc.agentFlagsStore.Update(kvPairs) // Extract just the keys from the key-value pairs - var updatedKeys []string - for i := 0; i < len(pairs); i += 2 { - updatedKeys = append(updatedKeys, pairs[i]) + updatedKeys := make([]string, len(kvPairs)) + i := 0 + for k := range kvPairs { + updatedKeys[i] = k + i++ } // Changed keys is the union of updated keys and deleted keys diff --git a/pkg/agent/flags/flag_controller_test.go b/pkg/agent/flags/flag_controller_test.go index c7c502633..9a432ff52 100644 --- a/pkg/agent/flags/flag_controller_test.go +++ b/pkg/agent/flags/flag_controller_test.go @@ -189,12 +189,12 @@ func TestControllerUpdate(t *testing.T) { tests := []struct { name string - pairs []string + kvPairs map[string]string changedKeys []string }{ { name: "happy path", - pairs: []string{keys.ControlRequestInterval.String(), "125000", keys.ControlServerURL.String(), "kolide-app.com"}, + kvPairs: map[string]string{keys.ControlRequestInterval.String(): "125000", keys.ControlServerURL.String(): "kolide-app.com"}, changedKeys: []string{keys.ControlRequestInterval.String(), keys.ControlServerURL.String()}, }, } @@ -213,7 +213,7 @@ func TestControllerUpdate(t *testing.T) { fc.RegisterChangeObserver(mockObserver, keys.ToFlagKeys(tt.changedKeys)...) - changedKeys, err := fc.Update(tt.pairs...) + changedKeys, err := fc.Update(tt.kvPairs) require.NoError(t, err) assert.Equal(t, tt.changedKeys, changedKeys) diff --git a/pkg/agent/storage/bbolt/keyvalue_store_bbolt.go b/pkg/agent/storage/bbolt/keyvalue_store_bbolt.go index 14959db3a..29c16d2dd 100644 --- a/pkg/agent/storage/bbolt/keyvalue_store_bbolt.go +++ b/pkg/agent/storage/bbolt/keyvalue_store_bbolt.go @@ -138,16 +138,11 @@ func (s *bboltKeyValueStore) ForEach(fn func(k, v []byte) error) error { }) } -func (s *bboltKeyValueStore) Update(pairs ...string) ([]string, error) { +func (s *bboltKeyValueStore) Update(kvPairs map[string]string) ([]string, error) { if s == nil || s.db == nil { return nil, NoDbError{} } - kvPairs := make(map[string]string) - for i := 0; i < len(pairs)-1; i += 2 { - kvPairs[pairs[i]] = pairs[i+1] - } - err := s.db.Update(func(tx *bbolt.Tx) error { b := tx.Bucket([]byte(s.bucketName)) if b == nil { diff --git a/pkg/agent/storage/ci/keyvalue_store_test.go b/pkg/agent/storage/ci/keyvalue_store_test.go index f24fa62eb..2961bbcbe 100644 --- a/pkg/agent/storage/ci/keyvalue_store_test.go +++ b/pkg/agent/storage/ci/keyvalue_store_test.go @@ -161,17 +161,17 @@ func Test_Updates(t *testing.T) { tests := []struct { name string - updates [][]string + updates []map[string]string want []map[string]string }{ { name: "empty", - updates: [][]string{{}, {}}, + updates: []map[string]string{{}, {}}, want: []map[string]string{}, }, { name: "single", - updates: [][]string{{"one", "one"}, {"one", "new_one"}}, + updates: []map[string]string{{"one": "one"}, {"one": "new_one"}}, want: []map[string]string{ { "key": "one", @@ -181,16 +181,16 @@ func Test_Updates(t *testing.T) { }, { name: "multiple", - updates: [][]string{ + updates: []map[string]string{ { - "one", "one", - "two", "two", - "three", "three", + "one": "one", + "two": "two", + "three": "three", }, { - "one", "new_one", - "two", "new_two", - "three", "new_three", + "one": "new_one", + "two": "new_two", + "three": "new_three", }, }, want: []map[string]string{ @@ -210,17 +210,17 @@ func Test_Updates(t *testing.T) { }, { name: "delete stale keys", - updates: [][]string{ + updates: []map[string]string{ { - "one", "one", - "two", "two", - "three", "three", - "four", "four", - "five", "five", - "six", "six", + "one": "one", + "two": "two", + "three": "three", + "four": "four", + "five": "five", + "six": "six", }, { - "four", "four", + "four": "four", }, }, want: []map[string]string{ @@ -238,7 +238,7 @@ func Test_Updates(t *testing.T) { for _, s := range getStores(t) { for _, update := range tt.updates { - s.Update(update...) + s.Update(update) } kvps, err := getKeyValueRows(s, tt.name) diff --git a/pkg/agent/storage/inmemory/keyvalue_store_in_memory.go b/pkg/agent/storage/inmemory/keyvalue_store_in_memory.go index 3d93d5981..5a3cf9d61 100644 --- a/pkg/agent/storage/inmemory/keyvalue_store_in_memory.go +++ b/pkg/agent/storage/inmemory/keyvalue_store_in_memory.go @@ -78,7 +78,7 @@ func (s *inMemoryKeyValueStore) ForEach(fn func(k, v []byte) error) error { return nil } -func (s *inMemoryKeyValueStore) Update(pairs ...string) ([]string, error) { +func (s *inMemoryKeyValueStore) Update(kvPairs map[string]string) ([]string, error) { if s == nil { return nil, errors.New("store is nil") } @@ -86,11 +86,6 @@ func (s *inMemoryKeyValueStore) Update(pairs ...string) ([]string, error) { s.mu.Lock() defer s.mu.Unlock() - kvPairs := make(map[string]string) - for i := 0; i < len(pairs)-1; i += 2 { - kvPairs[pairs[i]] = pairs[i+1] - } - s.items = make(map[string][]byte) for key, value := range kvPairs { diff --git a/pkg/agent/types/keyvalue_store.go b/pkg/agent/types/keyvalue_store.go index e6c1b761a..4616bed01 100644 --- a/pkg/agent/types/keyvalue_store.go +++ b/pkg/agent/types/keyvalue_store.go @@ -33,10 +33,10 @@ type Iterator interface { // Updater is an interface for bulk replacing data in a key/value store. type Updater interface { - // Update takes a sequence of alternating key-value pairs, and inserts + // Update takes a map of key-value pairs, and inserts // these key-values into the store. Any preexisting keys in the store which // do not exist in data will be deleted. - Update(pairs ...string) ([]string, error) + Update(kvPairs map[string]string) ([]string, error) } // GetterSetter is an interface that groups the Get and Set methods. diff --git a/pkg/agent/types/mocks/keyvalue_store.go b/pkg/agent/types/mocks/keyvalue_store.go index e9a2dcff1..2fda48617 100644 --- a/pkg/agent/types/mocks/keyvalue_store.go +++ b/pkg/agent/types/mocks/keyvalue_store.go @@ -83,31 +83,25 @@ func (_m *GetterSetterDeleterIteratorUpdater) Set(key []byte, value []byte) erro return r0 } -// Update provides a mock function with given fields: pairs -func (_m *GetterSetterDeleterIteratorUpdater) Update(pairs ...string) ([]string, error) { - _va := make([]interface{}, len(pairs)) - for _i := range pairs { - _va[_i] = pairs[_i] - } - var _ca []interface{} - _ca = append(_ca, _va...) - ret := _m.Called(_ca...) +// Update provides a mock function with given fields: kvPairs +func (_m *GetterSetterDeleterIteratorUpdater) Update(kvPairs map[string]string) ([]string, error) { + ret := _m.Called(kvPairs) var r0 []string var r1 error - if rf, ok := ret.Get(0).(func(...string) ([]string, error)); ok { - return rf(pairs...) + if rf, ok := ret.Get(0).(func(map[string]string) ([]string, error)); ok { + return rf(kvPairs) } - if rf, ok := ret.Get(0).(func(...string) []string); ok { - r0 = rf(pairs...) + if rf, ok := ret.Get(0).(func(map[string]string) []string); ok { + r0 = rf(kvPairs) } else { if ret.Get(0) != nil { r0 = ret.Get(0).([]string) } } - if rf, ok := ret.Get(1).(func(...string) error); ok { - r1 = rf(pairs...) + if rf, ok := ret.Get(1).(func(map[string]string) error); ok { + r1 = rf(kvPairs) } else { r1 = ret.Error(1) }