Added suport for multi-collection workloads and upserts

Author: daniel-percona

README.md

@@ -43,7 +43,7 @@ tar -xzvf plgm-darwin-arm64.tar.gz
 This project includes a `Makefile` to simplify building and packaging.
 
 ```bash
-git clone <repository-url>
+git clone https://github.com/Percona-Lab/percona-load-generator-mongodb.git
 cd plgm
 go mod tidy
 
@@ -182,8 +182,8 @@ You can override any setting in `config.yaml` using environment variables. This
 | `concurrency` | `PERCONALOAD_CONCURRENCY` | Number of active worker goroutines | `50` |
 | `duration` | `PERCONALOAD_DURATION` | Test duration (Go duration string) | `5m`, `60s` |
 | `default_workload` | `PERCONALOAD_DEFAULT_WORKLOAD` | Use built-in "Flights" workload (`true`/`false`) | `false` |
-| `collections_path` | `PERCONALOAD_COLLECTIONS_PATH` | Path to custom collection JSON files | `./schemas` |
-| `queries_path` | `PERCONALOAD_QUERIES_PATH` | Path to custom query JSON files | `./queries` |
+| `collections_path` | `PERCONALOAD_COLLECTIONS_PATH` | Path to custom collection JSON files (supports directories for multi-collection load) | `./schemas` |
+| `queries_path` | `PERCONALOAD_QUERIES_PATH` | Path to custom query JSON files or directory. | `./queries` |
 | `documents_count` | `PERCONALOAD_DOCUMENTS_COUNT` | Number of documents to seed initially | `10000` |
 | `drop_collections` | `PERCONALOAD_DROP_COLLECTIONS` | Drop collections before starting (`true`/`false`) | `true` |
 | `skip_seed` | `PERCONALOAD_SKIP_SEED` | Do not seed initial data on start (`true`/`false`) | `true` |
@@ -225,8 +225,9 @@ When executed, plgm performs the following steps:
     * (Optional) Seeds initial data with the number of documents defined by `documents_count` in the config.
 3.  **Workload Execution:**
     * Spawns the configured number of **Active Workers**.
-    * Continuously generates and executes queries (Find, Insert, Update, Delete, Aggregate) based on your configured ratios.
+    * Continuously generates and executes queries (Find, Insert, Update, Delete, Aggregate, Upsert) based on your configured ratios.
     * Generates realistic BSON data for Inserts and Updates (supports recursion and complex schemas).
+    * Workers pick a random collection from the provided list for every operation.
 4.  **Reporting:**
     * Outputs a real-time status report every N seconds (configurable).
     * Prints a detailed summary table at the end of the run.
@@ -323,6 +324,14 @@ To run your own workload against your own schema:
         "operation": "find",
         "filter": { "customer_name": "<string>" },
         "limit": 10
+      },
+      {
+        "database": "ecommerce",
+        "collection": "orders",
+        "operation": "updateOne",
+        "filter": { "order_uuid": "<string>" },
+        "update": { "$set": { "status": "processed" } },
+        "upsert": true
       }
     ]
     ```
@@ -383,6 +392,7 @@ Please note:
 * If there are no aggregation queries defined in queries.json, the aggregate_percent value is also ignored.
 * Aggregate operations will only generate activity if at least one query with "operation": "aggregate" is defined in your active JSON query files.
 * The maximum number of operations within a transaction is defined in the config file via `max_transaction_ops` or the env var `PERCONALOAD_MAX_TRANSACTION_OPS`. The number of operations per transaction will be randomized, with the max number being set as explained above. 
+* Multi-Collection Load: If multiple collections are defined in your collections_path, each worker will randomly select a collection for every operation. This includes operations within a transaction, allowing for cross-collection atomic updates.
 
 
 #### Concurrency & Workers
@@ -413,6 +423,7 @@ These settings affect the efficiency of individual database operations and memor
 * **`insert_cache_size`**: The buffer size for the document generator channel.
     * *Tip:* This decouples document generation from database insertion. A larger buffer ensures workers rarely wait for data generation logic.
     * *Default:* `1000`
+* **`Upserts`**: Any updateOne or updateMany operation in your query JSON files can include "upsert": true. This will cause MongoDB to create the document if no match is found for the filter.    
 
 #### Timeouts & Reliability
 Control how plgm reacts to network lag or database pressure.

internal/config/queries.go

@@ -17,6 +17,7 @@ type QueryDefinition struct {
 	Projection map[string]interface{} `json:"projection,omitempty" yaml:"projection,omitempty"`
 	Limit      int64                  `json:"limit,omitempty" yaml:"limit,omitempty"`
 	Update     map[string]interface{} `json:"update,omitempty" yaml:"update,omitempty"`
+	Upsert     bool                   `json:"upsert,omitempty" yaml:"upsert,omitempty"`
 }
 
 type QueriesFile struct {

internal/mongo/runner.go

@@ -171,14 +171,14 @@ func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.
 	}
 	defer session.EndSession(ctx)
 
-	mainCol := wCfg.collections[0]
 	start := time.Now()
 
 	_, err = session.WithTransaction(ctx, func(sessCtx context.Context) (interface{}, error) {
-		// generates a random number between 1 and MaxTransactionOps, this controls the number of random operations per transaction
 		numOps := rng.Intn(wCfg.appConfig.MaxTransactionOps) + 1
 		for i := 0; i < numOps; i++ {
-			// Select standard CRUD
+			// Select random collection for this step of the transaction
+			currentCol := wCfg.collections[rng.Intn(len(wCfg.collections))]
+
 			innerOp := selectOperation(wCfg.percentages, rng)
 			if innerOp == "aggregate" || innerOp == "transaction" {
 				innerOp = "find"
@@ -188,10 +188,10 @@ func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.
 			var run bool
 
 			if innerOp == "insert" {
-				q = generateInsertQuery(mainCol, rng, wCfg.appConfig)
+				q = generateInsertQuery(currentCol, rng, wCfg.appConfig)
 				run = true
 			} else {
-				q, run = selectRandomQueryByType(sessCtx, wCfg.database, innerOp, wCfg.queryMap, mainCol, wCfg.debug, rng, wCfg.primaryFilterField, wCfg.appConfig)
+				q, run = selectRandomQueryByType(sessCtx, wCfg.database, innerOp, wCfg.queryMap, currentCol, wCfg.debug, rng, wCfg.primaryFilterField, wCfg.appConfig)
 			}
 
 			if !run {
@@ -210,8 +210,12 @@ func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.
 					}
 					_ = cursor.Close(sessCtx)
 				}
-			case "updateOne", "updateMany":
-				_, err = coll.UpdateOne(sessCtx, filter, q.Update)
+			case "updateOne":
+				opts := options.UpdateOne().SetUpsert(q.Upsert)
+				_, err = coll.UpdateOne(sessCtx, filter, q.Update, opts)
+			case "updateMany":
+				opts := options.UpdateMany().SetUpsert(q.Upsert)
+				_, err = coll.UpdateMany(sessCtx, filter, q.Update, opts)
 			case "deleteOne", "deleteMany":
 				_, err = coll.DeleteOne(sessCtx, filter)
 			case "insert":
@@ -238,7 +242,6 @@ func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.
 func independentWorker(ctx context.Context, id int, wg *sync.WaitGroup, wCfg workloadConfig, rng *rand.Rand) {
 	defer wg.Done()
 	dbOpCtx := context.Background()
-	mainCol := wCfg.collections[0]
 
 	for {
 		select {
@@ -247,23 +250,28 @@ func independentWorker(ctx context.Context, id int, wg *sync.WaitGroup, wCfg wor
 		default:
 		}
 
+		// Pick random collection for this operation
+		currentCol := wCfg.collections[rng.Intn(len(wCfg.collections))]
+
 		opType := selectOperation(wCfg.percentages, rng)
 
-		// Handle Transaction Block
 		if opType == "transaction" {
-			runTransaction(ctx, id, wCfg, rng)
-			continue
+			if wCfg.appConfig.UseTransactions {
+				runTransaction(ctx, id, wCfg, rng)
+				continue
+			}
+			opType = "find"
 		}
 
 		var q config.QueryDefinition
 		var run bool
 
 		switch opType {
 		case "insert":
-			q = generateInsertQuery(mainCol, rng, wCfg.appConfig)
+			q = generateInsertQuery(currentCol, rng, wCfg.appConfig)
 			run = true
 		case "find", "updateOne", "updateMany", "deleteOne", "deleteMany", "aggregate":
-			q, run = selectRandomQueryByType(dbOpCtx, wCfg.database, opType, wCfg.queryMap, mainCol, wCfg.debug, rng, wCfg.primaryFilterField, wCfg.appConfig)
+			q, run = selectRandomQueryByType(dbOpCtx, wCfg.database, opType, wCfg.queryMap, currentCol, wCfg.debug, rng, wCfg.primaryFilterField, wCfg.appConfig)
 		default:
 			time.Sleep(100 * time.Microsecond)
 			continue
@@ -318,8 +326,20 @@ func independentWorker(ctx context.Context, id int, wg *sync.WaitGroup, wCfg wor
 				}
 				_ = cursor.Close(dbOpCtx)
 			}
-		case "updateOne", "updateMany":
-			coll.UpdateOne(dbOpCtx, filter, q.Update)
+		case "updateOne":
+			// Use UpdateOne() for single document updates in v2
+			opts := options.UpdateOne().SetUpsert(q.Upsert)
+			_, err := coll.UpdateOne(dbOpCtx, filter, q.Update, opts)
+			if err != nil && wCfg.debug {
+				log.Printf("[Worker %d] UpdateOne error: %v", id, err)
+			}
+		case "updateMany":
+			// Use UpdateMany() for multiple document updates in v2
+			opts := options.UpdateMany().SetUpsert(q.Upsert)
+			_, err := coll.UpdateMany(dbOpCtx, filter, q.Update, opts)
+			if err != nil && wCfg.debug {
+				log.Printf("[Worker %d] UpdateMany error: %v", id, err)
+			}
 		case "deleteOne", "deleteMany":
 			coll.DeleteOne(dbOpCtx, filter)
 		case "insert":
@@ -435,8 +455,10 @@ func runContinuousWorkload(ctx context.Context, wCfg workloadConfig) error {
 	workloadCtx, cancel := context.WithTimeout(ctx, wCfg.duration)
 	defer cancel()
 
-	mainCol := wCfg.collections[0]
-	go insertDocumentProducer(workloadCtx, mainCol, wCfg.maxInsertCache, wCfg.appConfig)
+	// Handle initial document producer for multiple collections
+	for _, col := range wCfg.collections {
+		go insertDocumentProducer(workloadCtx, col, wCfg.maxInsertCache, wCfg.appConfig)
+	}
 
 	monitorDone := make(chan struct{})
 	go func() {