diff --git a/splitio/proxy/storage/large_segments.go b/splitio/proxy/storage/large_segments.go new file mode 100644 index 00000000..974da9e0 --- /dev/null +++ b/splitio/proxy/storage/large_segments.go @@ -0,0 +1,67 @@ +package storage + +import ( + "sort" + "sync" + + "github.com/splitio/go-toolkit/v5/logging" +) + +// LargeSegmentsStorage defines the interface for a per-user large segments storage +type LargeSegmentsStorage interface { + Count() int + LargeSegmentsForUser(userKey string) []string + Update(lsName string, userKeys []string) +} + +// LargeSegmentsStorageImpl implements the LargeSegmentsStorage interface +type LargeSegmentsStorageImpl struct { + largeSegments map[string][]string + mutex *sync.RWMutex + logger logging.LoggerInterface +} + +// NewLargeSegmentsStorage constructs a new LargeSegments cache +func NewLargeSegmentsStorage(logger logging.LoggerInterface) *LargeSegmentsStorageImpl { + return &LargeSegmentsStorageImpl{ + largeSegments: make(map[string][]string), + mutex: &sync.RWMutex{}, + logger: logger, + } +} + +// Count retuns the amount of Large Segments +func (s *LargeSegmentsStorageImpl) Count() int { + s.mutex.RLock() + defer s.mutex.RUnlock() + return len(s.largeSegments) +} + +// SegmentsForUser returns the list of segments a certain user belongs to +func (s *LargeSegmentsStorageImpl) LargeSegmentsForUser(userKey string) []string { + s.mutex.RLock() + defer s.mutex.RUnlock() + + toReturn := make([]string, 0, len(s.largeSegments)) + for lsName, data := range s.largeSegments { + i := sort.Search(len(data), func(i int) bool { + return data[i] >= userKey + }) + + if i < len(data) && data[i] == userKey { + toReturn = append(toReturn, lsName) + } + } + + return toReturn +} + +// Update adds and remove keys to segments +func (s *LargeSegmentsStorageImpl) Update(lsName string, userKeys []string) { + s.mutex.Lock() + defer s.mutex.Unlock() + + s.largeSegments[lsName] = userKeys +} + +var _ LargeSegmentsStorage = (*LargeSegmentsStorageImpl)(nil) diff --git a/splitio/proxy/storage/large_segments_test.go b/splitio/proxy/storage/large_segments_test.go new file mode 100644 index 00000000..44546ccd --- /dev/null +++ b/splitio/proxy/storage/large_segments_test.go @@ -0,0 +1,56 @@ +package storage + +import ( + "sort" + "testing" + + "github.com/google/uuid" + "github.com/splitio/go-toolkit/v5/logging" + "github.com/stretchr/testify/assert" +) + +func sortedKeys(count int, shared *string) []string { + keys := make([]string, 0, count) + for i := 0; i < count; i++ { + keys = append(keys, uuid.New().String()) + } + + if shared != nil { + keys = append(keys, *shared) + } + + sort.Strings(keys) + return keys +} + +func TestLatgeSegmentStorage(t *testing.T) { + storage := NewLargeSegmentsStorage(logging.NewLogger(nil)) + + keys1 := sortedKeys(10000, nil) + storage.Update("ls_test_1", keys1) + + sharedKey := &keys1[5000] + keys2 := sortedKeys(20000, sharedKey) + storage.Update("ls_test_2", keys2) + + keys3 := sortedKeys(30000, sharedKey) + storage.Update("ls_test_3", keys3) + + assert.Equal(t, 3, storage.Count()) + + result := storage.LargeSegmentsForUser(*sharedKey) + sort.Strings(result) + assert.Equal(t, []string{"ls_test_1", "ls_test_2", "ls_test_3"}, result) + + result = storage.LargeSegmentsForUser(keys1[100]) + assert.Equal(t, []string{"ls_test_1"}, result) + + result = storage.LargeSegmentsForUser(keys2[100]) + assert.Equal(t, []string{"ls_test_2"}, result) + + result = storage.LargeSegmentsForUser(keys3[100]) + assert.Equal(t, []string{"ls_test_3"}, result) + + result = storage.LargeSegmentsForUser("mauro-test") + assert.Equal(t, []string{}, result) +}