From 2665b8829764c3961537fffcaa03292ea9074d64 Mon Sep 17 00:00:00 2001
From: Richard Lin <richard@rockset.com>
Date: Tue, 16 May 2023 18:51:53 -0700
Subject: [PATCH] Add functionality to generate docs with same field values
 over some range

---
 generator/document.go | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/generator/document.go b/generator/document.go
index 7ee8537..74f5839 100644
--- a/generator/document.go
+++ b/generator/document.go
@@ -109,12 +109,25 @@ func CurrentTimeMicros() int64 {
 func GenerateDocs(batchSize int, destination, identifier string, idMode string) ([]interface{}, error) {
 	var docs = make([]interface{}, batchSize, batchSize)
 
-	for i := 0; i < batchSize; i++ {
-		doc, err := GenerateDoc(destination, identifier, idMode)
-		if err != nil {
-			return nil, err
-		}
-		docs[i] = doc
+	i := 0
+	for i < batchSize {
+		// Choose between 10 and 20 documents to have the same field value
+		numSameField := rand.Intn(11) + 10
+		// Use uuid so that different kubernetes replicas won't generate the same values
+		claimedBy := guuid.New().String()
+		for j := 0; j < numSameField; j++ {
+			doc, err := GenerateDoc(destination, identifier, idMode)
+			if err != nil {
+				return nil, err
+			}
+			doc.(map[string]interface{})["claimedBy"] = claimedBy
+			docs[i] = doc
+			i++
+			// Break early once we reach the batch size
+			if (i >= batchSize) {
+				break
+			}
+		}		
 	}
 
 	return docs, nil