adding process implementation to handle unique keys by batches

Author: sanzmauro

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/gin-gonic/gin v1.10.0
 	github.com/google/uuid v1.3.0
 	github.com/splitio/gincache v1.0.1
-	github.com/splitio/go-split-commons/v6 v6.1.0
+	github.com/splitio/go-split-commons/v6 v6.1.1-0.20250414224340-9a5c36389db1
 	github.com/splitio/go-toolkit/v5 v5.4.0
 	github.com/stretchr/testify v1.9.0
 	go.etcd.io/bbolt v1.3.6

go.sum

@@ -93,8 +93,8 @@ github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUA
 github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
 github.com/splitio/gincache v1.0.1 h1:dLYdANY/BqH4KcUMCe/LluLyV5WtuE/LEdQWRE06IXU=
 github.com/splitio/gincache v1.0.1/go.mod h1:CcgJDSM9Af75kyBH0724v55URVwMBuSj5x1eCWIOECY=
-github.com/splitio/go-split-commons/v6 v6.1.0 h1:k3mwr12DF6gbEaV8XXU/tSAQlPkIEuzIgTEneYhGg2I=
-github.com/splitio/go-split-commons/v6 v6.1.0/go.mod h1:D/XIY/9Hmfk9ivWsRsJVp439kEdmHbzUi3PKzQQDOXY=
+github.com/splitio/go-split-commons/v6 v6.1.1-0.20250414224340-9a5c36389db1 h1:le2tNqtYsipMP5PQgJWAtu3x5DbPzquYT36RUeiiVfY=
+github.com/splitio/go-split-commons/v6 v6.1.1-0.20250414224340-9a5c36389db1/go.mod h1:D/XIY/9Hmfk9ivWsRsJVp439kEdmHbzUi3PKzQQDOXY=
 github.com/splitio/go-toolkit/v5 v5.4.0 h1:g5WFpRhQomnXCmvfsNOWV4s5AuUrWIZ+amM68G8NBKM=
 github.com/splitio/go-toolkit/v5 v5.4.0/go.mod h1:xYhUvV1gga9/1029Wbp5pjnR6Cy8nvBpjw99wAbsMko=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

splitio/producer/conf/sections.go

@@ -70,7 +70,7 @@ type AdvancedSync struct {
 	EventsPostConcurrency            int   `json:"eventsPostConcurrency" s-cli:"events-post-concurrency" s-def:"0" s-desc:"#concurrent imp post threads"`
 	EventsPostSize                   int   `json:"eventsPostSize" s-cli:"events-post-size" s-def:"0" s-desc:"Max #impressions to send per POST"`
 	EventsAccumWaitMs                int64 `json:"eventsAccumWaitMs" s-cli:"events-accum-wait-ms" s-def:"0" s-desc:"Max ms to wait to close an events bulk"`
-	UniqueKeysFetchSize              int64 `json:"uniqueKeysFetchSize" s-cli:"unique-keys-fetch-size" s-def:"0" s-desc:"How many unique keys to pop from storage at once"`
+	UniqueKeysFetchSize              int64 `json:"uniqueKeysFetchSize" s-cli:"unique-keys-fetch-size" s-def:"1000" s-desc:"How many unique keys to pop from storage at once"`
 	UniqueKeysProcessConcurrency     int   `json:"uniqueKeysProcessConcurrency" s-cli:"unique-keys-process-concurrency" s-def:"0" s-desc:"#Threads for processing uniques"`
 	UniqueKeysProcessBatchSize       int   `json:"uniqueKeysProcessBatchSize" s-cli:"unique-keys-process-batch-size" s-def:"0" s-desc:"Size of uniques processing batchs"`
 	UniqueKeysPostConcurrency        int   `json:"uniqueKeysPostConcurrency" s-cli:"unique-keys-post-concurrency" s-def:"0" s-desc:"#concurrent uniques post threads"`

splitio/producer/initialization.go

@@ -132,7 +132,7 @@ func Start(logger logging.LoggerInterface, cfg *conf.Main) error {
 			metadata, logger, syncTelemetryStorage),
 		// local telemetry
 		TelemetryRecorder: telemetry.NewTelemetrySynchronizer(syncTelemetryStorage, splitAPI.TelemetryRecorder,
-			storages.SplitStorage, storages.SegmentStorage, logger, metadata, syncTelemetryStorage),
+			storages.SplitStorage, storages.SegmentStorage, logger, metadata, syncTelemetryStorage, nil),
 	}
 	splitTasks := synchronizer.SplitTasks{
 		SplitSyncTask: tasks.NewFetchSplitsTask(workers.SplitUpdater, int(cfg.Sync.SplitRefreshRateMs)/1000, logger),
@@ -216,15 +216,14 @@ func Start(logger logging.LoggerInterface, cfg *conf.Main) error {
 	}
 
 	filter := filter.NewBloomFilter(bfExpectedElemenets, bfFalsePositiveProbability)
-	uniqueKeysTracker := strategy.NewUniqueKeysTracker(filter)
+
 	uniquesWorker := task.NewUniqueKeysWorker(&task.UniqueWorkerConfig{
-		Logger:            logger,
-		Storage:           storages.UniqueKeysStorage,
-		UniqueKeysTracker: uniqueKeysTracker,
-		URL:               advanced.TelemetryServiceURL,
-		Apikey:            cfg.Apikey,
-		FetchSize:         int(cfg.Sync.Advanced.UniqueKeysFetchSize),
-		Metadata:          metadata,
+		Logger:    logger,
+		Storage:   storages.UniqueKeysStorage,
+		URL:       advanced.TelemetryServiceURL,
+		Apikey:    cfg.Apikey,
+		FetchSize: int(cfg.Sync.Advanced.UniqueKeysFetchSize),
+		Metadata:  metadata,
 	})
 
 	uniquesTask, err := task.NewPipelinedTask(&task.Config{

splitio/producer/task/uniquekeys.go

@@ -7,48 +7,44 @@ import (
 	"net/http"
 
 	"github.com/splitio/go-split-commons/v6/dtos"
-	"github.com/splitio/go-split-commons/v6/provisional/strategy"
 	"github.com/splitio/go-split-commons/v6/storage"
+	"github.com/splitio/go-toolkit/v5/datastructures/set"
 	"github.com/splitio/go-toolkit/v5/logging"
 )
 
 // UniqueWorkerConfig bundles options
 type UniqueWorkerConfig struct {
-	Logger            logging.LoggerInterface
-	Storage           storage.UniqueKeysMultiSdkConsumer
-	UniqueKeysTracker strategy.UniqueKeysTracker
-	URL               string
-	Apikey            string
-	FetchSize         int
-	Metadata          dtos.Metadata
+	Logger    logging.LoggerInterface
+	Storage   storage.UniqueKeysMultiSdkConsumer
+	URL       string
+	Apikey    string
+	FetchSize int
+	Metadata  dtos.Metadata
 }
 
 // UniqueKeysPipelineWorker implements all the required  methods to work with a pipelined task
 type UniqueKeysPipelineWorker struct {
-	logger            logging.LoggerInterface
-	storage           storage.UniqueKeysMultiSdkConsumer
-	uniqueKeysTracker strategy.UniqueKeysTracker
-
+	logger    logging.LoggerInterface
+	storage   storage.UniqueKeysMultiSdkConsumer
 	url       string
 	apikey    string
-	fetchSize int64
+	fetchSize int
 	metadata  dtos.Metadata
 }
 
 func NewUniqueKeysWorker(cfg *UniqueWorkerConfig) Worker {
 	return &UniqueKeysPipelineWorker{
-		logger:            cfg.Logger,
-		storage:           cfg.Storage,
-		uniqueKeysTracker: cfg.UniqueKeysTracker,
-		url:               cfg.URL + "/keys/ss",
-		apikey:            cfg.Apikey,
-		fetchSize:         int64(cfg.FetchSize),
-		metadata:          cfg.Metadata,
+		logger:    cfg.Logger,
+		storage:   cfg.Storage,
+		url:       cfg.URL + "/keys/ss",
+		apikey:    cfg.Apikey,
+		fetchSize: cfg.FetchSize,
+		metadata:  cfg.Metadata,
 	}
 }
 
 func (u *UniqueKeysPipelineWorker) Fetch() ([]string, error) {
-	raw, _, err := u.storage.PopNRaw(u.fetchSize)
+	raw, _, err := u.storage.PopNRaw(int64(u.fetchSize))
 	if err != nil {
 		return nil, fmt.Errorf("error fetching raw unique keys: %w", err)
 	}
@@ -57,31 +53,30 @@ func (u *UniqueKeysPipelineWorker) Fetch() ([]string, error) {
 }
 
 func (u *UniqueKeysPipelineWorker) Process(raws [][]byte, sink chan<- interface{}) error {
+	rawKeys := make([]dtos.Key, 0)
 	for _, raw := range raws {
-		err, value := parseToObj(raw)
+		value, err := parseToObj(raw)
 		if err == nil {
 			u.logger.Debug("Unique Keys parsed to Dto.")
 		}
 
 		if err != nil {
-			err, value = parseToArray(raw)
+			value, err = parseToArray(raw)
 			if err != nil {
 				u.logger.Error("error deserializing fetched uniqueKeys: ", err.Error())
 				continue
 			}
 			u.logger.Debug("Unique Keys parsed to Array.")
 		}
 
-		for _, unique := range value {
-			for _, key := range unique.Keys {
-				u.uniqueKeysTracker.Track(unique.Feature, key)
-			}
-		}
+		rawKeys = append(rawKeys, value...)
 	}
 
-	uniques := u.uniqueKeysTracker.PopAll()
-	if len(uniques.Keys) > 0 {
-		sink <- uniques
+	filtered := cleanUp(rawKeys)
+	groups := batches(filtered, u.fetchSize)
+
+	for index := range groups {
+		sink <- groups[index]
 	}
 
 	return nil
@@ -108,22 +103,78 @@ func (u *UniqueKeysPipelineWorker) BuildRequest(data interface{}) (*http.Request
 	return req, nil
 }
 
-func parseToArray(raw []byte) (error, []dtos.Key) {
+func parseToArray(raw []byte) ([]dtos.Key, error) {
 	var queueObj []dtos.Key
 	err := json.Unmarshal(raw, &queueObj)
 	if err != nil {
-		return err, nil
+		return nil, err
 	}
 
-	return nil, queueObj
+	return queueObj, nil
 }
 
-func parseToObj(raw []byte) (error, []dtos.Key) {
+func parseToObj(raw []byte) ([]dtos.Key, error) {
 	var queueObj dtos.Key
 	err := json.Unmarshal(raw, &queueObj)
 	if err != nil {
-		return err, nil
+		return nil, err
+	}
+
+	return []dtos.Key{queueObj}, nil
+}
+
+func cleanUp(keys []dtos.Key) map[string]*set.ThreadUnsafeSet {
+	filtered := make(map[string]*set.ThreadUnsafeSet)
+	for _, key := range keys {
+		_, exists := filtered[key.Feature]
+		if !exists {
+			filtered[key.Feature] = set.NewSet()
+		}
+		filtered[key.Feature].Add(key.Keys...)
+	}
+
+	return filtered
+}
+
+func batches(filtered map[string]*set.ThreadUnsafeSet, maxSize int) []dtos.Uniques {
+	groups := make([]dtos.Uniques, 0)
+	currentBatch := dtos.Uniques{Keys: []dtos.Key{}}
+	currentBatchSize := 0
+
+	for name, keys := range filtered {
+		keyList := keys.List()
+		totalKeys := len(keyList)
+		start := 0
+
+		for start < totalKeys {
+			end := start + maxSize - currentBatchSize
+			if end > totalKeys {
+				end = totalKeys
+			}
+
+			// Add keys to the current batch
+			keyDto := dtos.Key{
+				Feature: name,
+				Keys:    keyList[start:end],
+			}
+			currentBatch.Keys = append(currentBatch.Keys, keyDto)
+			currentBatchSize += len(keyList[start:end])
+
+			// If the current batch reaches maxSize, finalize it and start a new one
+			if currentBatchSize >= maxSize {
+				groups = append(groups, currentBatch)
+				currentBatch = dtos.Uniques{Keys: []dtos.Key{}}
+				currentBatchSize = 0
+			}
+
+			start = end
+		}
+	}
+
+	// Add the remaining batch if it has any keys
+	if currentBatchSize > 0 {
+		groups = append(groups, currentBatch)
 	}
 
-	return nil, []dtos.Key{queueObj}
+	return groups
 }