fix: Add ChangeSet.Collections to help reduce relay memory usage (#332)

The SDK changes to convert changesets into `[]ldstoretypes.Collection`
for a lot of it's internal workings. Coincidentally, the relay proxy
also needs to do this same conversion when supporting FDv1 formats.

To help with this, we add a `.Collections` method which memoizes the
result for later use, reducing our memory allocation.

Author: keelerm84

internal/datasourcev2/streaming_data_source_test.go

@@ -563,7 +563,7 @@ func TestStreamingDataSourceHandlesUpToDateAndSubsequentChanges(t *testing.T) {
 			status = <-resultChan
 			assert.Equal(t, interfaces.DataSourceStateValid, status.State)
 			assert.NotNil(t, status.ChangeSet)
-			dd.SetBasis(status.ChangeSet.Changes(), status.ChangeSet.Selector(), true)
+			dd.Apply(*status.ChangeSet, true)
 
 			flag, err = dd.DataStore.Get(datakinds.Features, "test-flag")
 			assert.NoError(t, err)
@@ -646,7 +646,7 @@ func TestStreamingDataSourceHandlesResettingFromError(t *testing.T) {
 			status = <-resultChan
 			assert.Equal(t, interfaces.DataSourceStateValid, status.State)
 			assert.NotNil(t, status.ChangeSet)
-			dd.SetBasis(status.ChangeSet.Changes(), status.ChangeSet.Selector(), true)
+			dd.Apply(*status.ChangeSet, true)
 
 			flag, err = dd.DataStore.Get(datakinds.Features, "test-flag")
 			assert.NoError(t, err)
@@ -734,7 +734,7 @@ func TestStreamingDataSourceIgnoresGoodbye(t *testing.T) {
 			status = <-resultChan
 			assert.Equal(t, interfaces.DataSourceStateValid, status.State)
 			assert.NotNil(t, status.ChangeSet)
-			dd.SetBasis(status.ChangeSet.Changes(), status.ChangeSet.Selector(), true)
+			dd.Apply(*status.ChangeSet, true)
 
 			flag, err = dd.DataStore.Get(datakinds.Features, "test-flag")
 			assert.NoError(t, err)

internal/datasystem/store.go

@@ -173,7 +173,7 @@ func (s *Store) Close() error {
 // Apply applies a changeset to the store. The changeset must be a valid set of changes that can be applied
 // to the store. If the changeset is not valid, an error will be logged and the changeset will not be applied.
 func (s *Store) Apply(changeSet subsystems.ChangeSet, persist bool) {
-	collections, err := subsystems.ToStorableItems(changeSet.Changes())
+	collections, err := changeSet.Collections()
 	if err != nil {
 		s.loggers.Errorf("store: couldn't set basis due to malformed data: %v", err)
 		return
@@ -244,7 +244,7 @@ func (s *Store) shouldPersist() bool {
 	return s.persist && s.persistentStore.writable()
 }
 
-// applyDelta applies a delta update to the store. applyDelta should not be called until SetBasis has been called.
+// applyDelta applies a delta update to the store. applyDelta should not be called until setBasis has been called.
 // To request data persistence, set persist to true.
 func (s *Store) applyDelta(collections []ldstoretypes.Collection, selector subsystems.Selector, persist bool) {
 	s.memoryStore.ApplyDelta(collections)
@@ -298,7 +298,7 @@ func (s *Store) GetDataStoreStatusProvider() interfaces.DataStoreStatusProvider
 }
 
 // Commit persists the data in the memory store to the persistent store, if configured. The persistent store
-// must also be in write mode, and the last call to SetBasis or ApplyDelta must have had persist set to true.
+// must also be in write mode, and the last call to setBasis or applyDelta must have had persist set to true.
 func (s *Store) Commit() error {
 	s.mu.RLock()
 	defer s.mu.RUnlock()

internal/sharedtest/mocks/mock_data_destination.go

@@ -50,14 +50,21 @@ func (d *MockDataDestination) Selector() subsystems.Selector {
 	return d.lastKnownSelector
 }
 
-// SetBasis in this test implementation, delegates to d.DataStore.CapturedUpdates.
-func (d *MockDataDestination) SetBasis(events []subsystems.Change, selector subsystems.Selector, _ bool) {
-	// For now, the selector is ignored. When the data sources start making use of it, it should be
-	// stored so that assertions can be made.
+// Apply persists the given ChangeSet to the DataStore.
+func (d *MockDataDestination) Apply(changeSet subsystems.ChangeSet, persist bool) {
+	switch changeSet.IntentCode() {
+	case subsystems.IntentTransferFull:
+		d.setBasis(changeSet, persist)
+	case subsystems.IntentTransferChanges:
+		d.applyDelta(changeSet, persist)
+	}
+}
 
-	collections, err := subsystems.ToStorableItems(events)
+// setBasis in this test implementation, delegates to d.DataStore.CapturedUpdates.
+func (d *MockDataDestination) setBasis(changeSet subsystems.ChangeSet, _ bool) {
+	collections, err := changeSet.Collections()
 	if err != nil {
-		panic("MockDataDestination.SetBasis received malformed data: " + err.Error())
+		panic("MockDataDestination.setBasis received malformed data: " + err.Error())
 	}
 
 	for _, coll := range collections {
@@ -66,19 +73,16 @@ func (d *MockDataDestination) SetBasis(events []subsystems.Change, selector subs
 
 	if err := d.DataStore.Init(toposort.Sort(collections)); err == nil {
 		d.lock.Lock()
-		d.lastKnownSelector = selector
+		d.lastKnownSelector = changeSet.Selector()
 		d.lock.Unlock()
 	}
 }
 
-// ApplyDelta in this test implementation, delegates to d.DataStore.CapturedUpdates.
-func (d *MockDataDestination) ApplyDelta(events []subsystems.Change, selector subsystems.Selector, _ bool) {
-	// For now, the selector is ignored. When the data sources start making use of it, it should be
-	// stored so that assertions can be made.
-
-	collections, err := subsystems.ToStorableItems(events)
+// applyDelta in this test implementation, delegates to d.DataStore.CapturedUpdates.
+func (d *MockDataDestination) applyDelta(changeSet subsystems.ChangeSet, _ bool) {
+	collections, err := changeSet.Collections()
 	if err != nil {
-		panic("MockDataDestination.ApplyDelta received malformed data: " + err.Error())
+		panic("MockDataDestination.applyDelta received malformed data: " + err.Error())
 	}
 
 	for _, coll := range collections {
@@ -94,7 +98,7 @@ func (d *MockDataDestination) ApplyDelta(events []subsystems.Change, selector su
 	}
 
 	d.lock.Lock()
-	d.lastKnownSelector = selector
+	d.lastKnownSelector = changeSet.Selector()
 	d.lock.Unlock()
 }
 

ldfiledatav2/file_data_source_impl.go

@@ -138,9 +138,11 @@ func (fs *fileDataSource) Fetch(ds subsystems.DataSelector, ctx context.Context)
 	changeSetChan := fs.changeSetBroadcaster.AddListener()
 	statusChan := fs.statusBroadcaster.AddListener()
 
+	changeset := subsystems.NewChangeSetBuilder().NoChanges()
+
 	var err error
 	basis := &subsystems.Basis{
-		ChangeSet: subsystems.ChangeSet{},
+		ChangeSet: *changeset,
 		Persist:   false,
 	}
 

subsystems/changeset.go

@@ -3,6 +3,10 @@ package subsystems
 import (
 	"encoding/json"
 	"errors"
+	"sync"
+
+	"github.com/launchdarkly/go-jsonstream/v3/jreader"
+	"github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoretypes"
 )
 
 // ChangeType specifies if an object is being upserted or deleted.
@@ -48,6 +52,9 @@ type ChangeSet struct {
 	intentCode IntentCode
 	changes    []Change
 	selector   Selector
+
+	mu         *sync.Mutex
+	collection []ldstoretypes.Collection
 }
 
 // IntentCode represents the intent of the changeset.
@@ -66,6 +73,70 @@ func (c *ChangeSet) Selector() Selector {
 	return c.selector
 }
 
+// Collections converts the changeset into a list of collections suitable for
+// insertion into a data store, relaying to FDv1 clients, etc.
+func (c *ChangeSet) Collections() ([]ldstoretypes.Collection, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.collection != nil {
+		return c.collection, nil
+	}
+
+	if c.changes == nil {
+		c.collection = []ldstoretypes.Collection{}
+		return c.collection, nil
+	}
+
+	collection, err := toStorableItems(c.changes)
+	if err != nil {
+		return nil, err
+	}
+
+	c.collection = collection
+	return c.collection, nil
+}
+
+// toStorableItems converts a list of FDv2 events to a list of collections suitable for insertion
+// into a data store.
+func toStorableItems(deltas []Change) ([]ldstoretypes.Collection, error) {
+	collections := make(kindMap)
+	for _, event := range deltas {
+		kind, ok := event.Kind.ToFDV1()
+		if !ok {
+			// If we don't recognize this kind, it's not an error and should be ignored for forwards
+			// compatibility.
+			continue
+		}
+
+		switch event.Action {
+		case ChangeTypePut:
+			// A put requires deserializing the item. We delegate to the optimized streaming JSON
+			// parser.
+			reader := jreader.NewReader(event.Object)
+			item, err := kind.DeserializeFromJSONReader(&reader)
+			if err != nil {
+				return nil, err
+			}
+			collections[kind] = append(collections[kind], ldstoretypes.KeyedItemDescriptor{
+				Key:  event.Key,
+				Item: item,
+			})
+		case ChangeTypeDelete:
+			// A deletion is represented by a tombstone, which is an ItemDescriptor with a version and nil item.
+			collections[kind] = append(collections[kind], ldstoretypes.KeyedItemDescriptor{
+				Key:  event.Key,
+				Item: ldstoretypes.ItemDescriptor{Version: event.Version, Item: nil},
+			})
+		default:
+			// An unknown action isn't an error, and should be ignored for forwards compatibility.
+			continue
+		}
+	}
+
+	return collections.flatten(), nil
+}
+
 // ChangeSetBuilder is a helper for constructing a ChangeSet.
 //
 // This type is not stable, and not subject to any backwards
@@ -96,6 +167,7 @@ func (c *ChangeSetBuilder) NoChanges() *ChangeSet {
 		intentCode: IntentNone,
 		selector:   NoSelector(),
 		changes:    nil,
+		mu:         &sync.Mutex{},
 	}
 }
 
@@ -106,6 +178,7 @@ func (c *ChangeSetBuilder) Empty(selector Selector) *ChangeSet {
 		intentCode: IntentTransferFull,
 		selector:   selector,
 		changes:    nil,
+		mu:         &sync.Mutex{},
 	}
 }
 
@@ -152,6 +225,7 @@ func (c *ChangeSetBuilder) Finish(selector Selector) (*ChangeSet, error) {
 		intentCode: c.intent.Payload.Code,
 		selector:   selector,
 		changes:    c.changes,
+		mu:         &sync.Mutex{},
 	}
 	c.changes = nil