fix(cmd): stop generators as last to avoid deadlock on exit

cmd/gemini/root.go

@@ -259,7 +259,8 @@ func run(_ *cobra.Command, _ []string) error {
 	stop.StartOsSignalsTransmitter(logger, &warmupStopFlag, &workStopFlag)
 	pump := jobs.NewPump(ctx, logger)
 
-	generators := createGenerators(ctx, schema, schemaConfig, distFunc, concurrency, partitionCount, logger)
+	genCtx, genCancel := context.WithCancel(context.Background())
+	generators := createGenerators(genCtx, schema, schemaConfig, distFunc, concurrency, partitionCount, logger)
 
 	if !nonInteractive {
 		sp := createSpinner(interactive())
@@ -294,8 +295,13 @@ func run(_ *cobra.Command, _ []string) error {
 		if err = jobsList.Run(ctx, schema, schemaConfig, st, pump, generators, globalStatus, logger, seed, &workStopFlag, failFast, verbose); err != nil {
 			logger.Debug("error detected", zap.Error(err))
 		}
-
 	}
+
+	// stop generators, this needs to be done after jobs finish
+	// because otherwise they can deadlock in pick() func waiting
+	// for more values to be generated
+	genCancel()
+
 	logger.Info("test finished")
 	globalStatus.PrintResult(outFile, schema, version)
 	if globalStatus.HasErrors() {

pkg/generators/generator.go

@@ -162,7 +162,7 @@ func (g *Generator) start() {
 			g.fillAllPartitions()
 			select {
 			case <-gCtx.Done():
-				g.logger.Debug("stopping partition key generation loop",
+				g.logger.Info("stopping partition key generation loop",
 					zap.Uint64("keys_created", g.cntCreated),
 					zap.Uint64("keys_emitted", g.cntEmitted))
 				return gCtx.Err()

pkg/jobs/jobs.go

@@ -177,7 +177,7 @@ func mutationJob(
 		}
 		select {
 		case <-ctx.Done():
-			logger.Debug("mutation job terminated")
+			logger.Info("mutation job terminated")
 			return nil
 		case hb := <-pump:
 			time.Sleep(hb)
@@ -286,7 +286,7 @@ func warmupJob(
 		}
 		select {
 		case <-ctx.Done():
-			logger.Debug("warmup job terminated")
+			logger.Info("warmup job terminated")
 			return nil
 		default:
 		}