refactor(stopFlag): completly remove stopFlag for context

`context` is Go package used mainly to signal cancelation. `stopFlag` package inside of gemini was used the same way. It was just adding overhead and maintainability issues in the long run, as it was only a wrapper for `context.Context` and all at once calling the `cancel` funcs at the end of the program. This removal makes code easier to follow for everybody working on it. `context` package is well known in Go community and can serve many purposes beyond just canceling the background workers (`jobs` in gemini terminology). Hard kill in gemini was used to signal the immadiet stoppage for gemini, without pritting the results, and soft on the contrary could be triggers in a few ways: - gemini validation failed - gemini mutation failed - SIG(INT,TERM) was sent There is no reason to have HARD kill in the application, if we need hard kill, SIGKILL can be sent and everything would be stopped, there will be no need to the the cleanup as kernel will ensure that happens. `context.Context` works as a soft kill, and if something happens bad in gemini (validation fails or mutation fail) `globalStatus.HasError()` would stop all other `goroutines` from continuing if `failFast` CLI flag is passed, so the set of `softKill` is not needed. Signed-off-by: Dusan Malusev <[email protected]>
scylladb · Dec 3, 2024 · 5f374f2 · 5f374f2
1 parent 69be9e9
commit 5f374f2
Show file tree

Hide file tree

Showing 8 changed files with 65 additions and 726 deletions.
diff --git a/cmd/gemini/root.go b/cmd/gemini/root.go
@@ -22,24 +22,13 @@ import (
 	"net/http"
 	"net/http/pprof"
 	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
 	"text/tabwriter"
 	"time"
 
-	"github.com/scylladb/gemini/pkg/auth"
-	"github.com/scylladb/gemini/pkg/builders"
-	"github.com/scylladb/gemini/pkg/generators"
-	"github.com/scylladb/gemini/pkg/jobs"
-	"github.com/scylladb/gemini/pkg/realrandom"
-	"github.com/scylladb/gemini/pkg/replication"
-	"github.com/scylladb/gemini/pkg/store"
-	"github.com/scylladb/gemini/pkg/typedef"
-	"github.com/scylladb/gemini/pkg/utils"
-
-	"github.com/scylladb/gemini/pkg/status"
-	"github.com/scylladb/gemini/pkg/stop"
-
 	"github.com/gocql/gocql"
 	"github.com/hailocab/go-hostpool"
 	"github.com/pkg/errors"
@@ -50,6 +39,17 @@ import (
 	"golang.org/x/exp/rand"
 	"golang.org/x/net/context"
 	"gonum.org/v1/gonum/stat/distuv"
+
+	"github.com/scylladb/gemini/pkg/auth"
+	"github.com/scylladb/gemini/pkg/builders"
+	"github.com/scylladb/gemini/pkg/generators"
+	"github.com/scylladb/gemini/pkg/jobs"
+	"github.com/scylladb/gemini/pkg/realrandom"
+	"github.com/scylladb/gemini/pkg/replication"
+	"github.com/scylladb/gemini/pkg/status"
+	"github.com/scylladb/gemini/pkg/store"
+	"github.com/scylladb/gemini/pkg/typedef"
+	"github.com/scylladb/gemini/pkg/utils"
 )
 
 var (
@@ -137,8 +137,11 @@ func readSchema(confFile string, schemaConfig typedef.SchemaConfig) (*typedef.Sc
 }
 
 func run(_ *cobra.Command, _ []string) error {
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGABRT, syscall.SIGTERM, syscall.SIGINT)
+	defer cancel()
+
 	logger := createLogger(level)
-	globalStatus := status.NewGlobalStatus(1000)
+	globalStatus := status.NewGlobalStatus(int32(maxErrorsToStore))
 	defer utils.IgnoreError(logger.Sync)
 
 	if err := validateSeed(seed); err != nil {
@@ -242,15 +245,15 @@ func run(_ *cobra.Command, _ []string) error {
 	if dropSchema && mode != jobs.ReadMode {
 		for _, stmt := range generators.GetDropKeyspace(schema) {
 			logger.Debug(stmt)
-			if err = st.Mutate(context.Background(), typedef.SimpleStmt(stmt, typedef.DropKeyspaceStatementType)); err != nil {
+			if err = st.Mutate(ctx, typedef.SimpleStmt(stmt, typedef.DropKeyspaceStatementType)); err != nil {
 				return errors.Wrap(err, "unable to drop schema")
 			}
 		}
 	}
 
 	testKeyspace, oracleKeyspace := generators.GetCreateKeyspaces(schema)
 	if err = st.Create(
-		context.Background(),
+		ctx,
 		typedef.SimpleStmt(testKeyspace, typedef.CreateKeyspaceStatementType),
 		typedef.SimpleStmt(oracleKeyspace, typedef.CreateKeyspaceStatementType)); err != nil {
 		return errors.Wrap(err, "unable to create keyspace")
@@ -263,11 +266,7 @@ func run(_ *cobra.Command, _ []string) error {
 		}
 	}
 
-	ctx, done := context.WithTimeout(context.Background(), duration+warmup+time.Second*2)
-	stopFlag := stop.NewFlag("main")
-	warmupStopFlag := stop.NewFlag("warmup")
-	stop.StartOsSignalsTransmitter(logger, stopFlag, warmupStopFlag)
-	pump := jobs.NewPump(stopFlag, logger)
+	pump := jobs.NewPump(ctx, logger)
 
 	distFunc, err := createDistributionFunc(partitionKeyDistribution, partitionCount, intSeed, normalDistMean, normalDistSigma)
 	if err != nil {
@@ -281,10 +280,9 @@ func run(_ *cobra.Command, _ []string) error {
 		sp := createSpinner(interactive())
 		ticker := time.NewTicker(time.Second)
 		go func() {
-			defer done()
 			for {
 				select {
-				case <-stopFlag.SignalChannel():
+				case <-ctx.Done():
 					return
 				case <-ticker.C:
 					sp.Set(" Running Gemini... %v", globalStatus)
@@ -293,20 +291,24 @@ func run(_ *cobra.Command, _ []string) error {
 		}()
 	}
 
-	if warmup > 0 && !stopFlag.IsHardOrSoft() {
-		jobsList := jobs.ListFromMode(jobs.WarmupMode, warmup, concurrency)
-		if err = jobsList.Run(ctx, schema, schemaConfig, st, pump, gens, globalStatus, logger, intSeed, warmupStopFlag, failFast, verbose); err != nil {
+	if warmup > 0 {
+		warmupCtx, warmupCancel := context.WithTimeout(ctx, warmup)
+		defer warmupCancel()
+
+		jobsList := jobs.ListFromMode(jobs.WarmupMode, concurrency)
+		if err = jobsList.Run(warmupCtx, schema, schemaConfig, st, pump, gens, globalStatus, logger, intSeed, failFast, verbose); err != nil {
 			logger.Error("warmup encountered an error", zap.Error(err))
-			stopFlag.SetHard(true)
 		}
 	}
 
-	if !stopFlag.IsHardOrSoft() {
-		jobsList := jobs.ListFromMode(mode, duration, concurrency)
-		if err = jobsList.Run(ctx, schema, schemaConfig, st, pump, gens, globalStatus, logger, intSeed, stopFlag.CreateChild("workload"), failFast, verbose); err != nil {
-			logger.Debug("error detected", zap.Error(err))
-		}
+	jobsCtx, jobsCancel := context.WithTimeout(ctx, duration)
+	defer jobsCancel()
+
+	jobsList := jobs.ListFromMode(mode, concurrency)
+	if err = jobsList.Run(jobsCtx, schema, schemaConfig, st, pump, gens, globalStatus, logger, intSeed, failFast, verbose); err != nil {
+		logger.Debug("error detected", zap.Error(err))
 	}
+
 	logger.Info("test finished")
 	globalStatus.PrintResult(outFile, schema, version)
 	if globalStatus.HasErrors() {

diff --git a/pkg/generators/generator.go b/pkg/generators/generator.go
@@ -27,7 +27,7 @@ import (
 
 // TokenIndex represents the position of a token in the token ring.
 // A token index is translated to a token by a generators. If the generators
-// preserves the exact position, then the token index becomes the token;
+// preserve the exact position, then the token index becomes the token;
 // otherwise token index represents an approximation of the token.
 //
 // We use a token index approach, because our generators actually generate

diff --git a/pkg/generators/generators.go b/pkg/generators/generators.go
@@ -25,9 +25,9 @@ import (
 )
 
 type Generators struct {
-	Generators []Generator
 	wg         *sync.WaitGroup
 	cancel     context.CancelFunc
+	Generators []Generator
 }
 
 func New(

diff --git a/pkg/generators/partition.go b/pkg/generators/partition.go
@@ -114,12 +114,11 @@ func (s *Partition) safelyGetOldValuesChannel() chan *typedef.ValueWithToken {
 }
 
 func (s *Partition) Close() error {
-	for !s.closed.CompareAndSwap(false, true) {
+	if s.closed.CompareAndSwap(false, true) {
+		close(s.values)
+		close(s.oldValues)
 	}
 
-	close(s.values)
-	close(s.oldValues)
-
 	return nil
 }
 

diff --git a/pkg/jobs/jobs.go b/pkg/jobs/jobs.go
@@ -28,7 +28,6 @@ import (
 	"github.com/scylladb/gemini/pkg/generators"
 	"github.com/scylladb/gemini/pkg/joberror"
 	"github.com/scylladb/gemini/pkg/status"
-	"github.com/scylladb/gemini/pkg/stop"
 	"github.com/scylladb/gemini/pkg/store"
 	"github.com/scylladb/gemini/pkg/typedef"
 )
@@ -53,10 +52,9 @@ var (
 )
 
 type List struct {
-	name     string
-	jobs     []job
-	duration time.Duration
-	workers  uint64
+	name    string
+	jobs    []job
+	workers uint64
 }
 
 type job struct {
@@ -72,16 +70,16 @@ type job struct {
 		*generators.Generator,
 		*status.GlobalStatus,
 		*zap.Logger,
-		*stop.Flag,
 		bool,
 		bool,
 	) error
 	name string
 }
 
-func ListFromMode(mode string, duration time.Duration, workers uint64) List {
+func ListFromMode(mode string, workers uint64) List {
 	jobs := make([]job, 0, 2)
 	name := "work cycle"
+
 	switch mode {
 	case WriteMode:
 		jobs = append(jobs, mutate)
@@ -93,11 +91,11 @@ func ListFromMode(mode string, duration time.Duration, workers uint64) List {
 	default:
 		jobs = append(jobs, mutate, validate)
 	}
+
 	return List{
-		name:     name,
-		jobs:     jobs,
-		duration: duration,
-		workers:  workers,
+		name:    name,
+		jobs:    jobs,
+		workers: workers,
 	}
 }
 
@@ -111,16 +109,10 @@ func (l List) Run(
 	globalStatus *status.GlobalStatus,
 	logger *zap.Logger,
 	seed uint64,
-	stopFlag *stop.Flag,
 	failFast, verbose bool,
 ) error {
 	logger = logger.Named(l.name)
-	ctx = stopFlag.CancelContextOnSignal(ctx, stop.SignalHardStop)
 	g, gCtx := errgroup.WithContext(ctx)
-	time.AfterFunc(l.duration, func() {
-		logger.Info("jobs time is up, begins jobs completion")
-		stopFlag.SetSoft(true)
-	})
 
 	partitionRangeConfig := schemaConfig.GetPartitionRangeConfig()
 	logger.Info("start jobs")
@@ -131,7 +123,7 @@ func (l List) Run(
 				jobF := l.jobs[idx].function
 				r := rand.New(rand.NewSource(seed))
 				g.Go(func() error {
-					return jobF(gCtx, pump, schema, schemaConfig, table, s, r, &partitionRangeConfig, generator, globalStatus, logger, stopFlag, failFast, verbose)
+					return jobF(gCtx, pump, schema, schemaConfig, table, s, r, &partitionRangeConfig, generator, globalStatus, logger, failFast, verbose)
 				})
 			}
 		}
@@ -154,7 +146,6 @@ func mutationJob(
 	g *generators.Generator,
 	globalStatus *status.GlobalStatus,
 	logger *zap.Logger,
-	stopFlag *stop.Flag,
 	failFast, verbose bool,
 ) error {
 	schemaConfig := &schemaCfg
@@ -164,11 +155,8 @@ func mutationJob(
 		logger.Info("ending mutation loop")
 	}()
 	for {
-		if stopFlag.IsHardOrSoft() {
-			return nil
-		}
 		select {
-		case <-stopFlag.SignalChannel():
+		case <-ctx.Done():
 			logger.Debug("mutation job terminated")
 			return nil
 		case hb := <-pump:
@@ -187,7 +175,6 @@ func mutationJob(
 			}
 		}
 		if failFast && globalStatus.HasErrors() {
-			stopFlag.SetSoft(true)
 			return nil
 		}
 	}
@@ -207,7 +194,6 @@ func validationJob(
 	g *generators.Generator,
 	globalStatus *status.GlobalStatus,
 	logger *zap.Logger,
-	stopFlag *stop.Flag,
 	failFast, _ bool,
 ) error {
 	schemaConfig := &schemaCfg
@@ -218,11 +204,8 @@ func validationJob(
 	}()
 
 	for {
-		if stopFlag.IsHardOrSoft() {
-			return nil
-		}
 		select {
-		case <-stopFlag.SignalChannel():
+		case <-ctx.Done():
 			return nil
 		case hb := <-pump:
 			time.Sleep(hb)
@@ -262,7 +245,6 @@ func validationJob(
 		}
 
 		if failFast && globalStatus.HasErrors() {
-			stopFlag.SetSoft(true)
 			return nil
 		}
 	}
@@ -282,7 +264,6 @@ func warmupJob(
 	g *generators.Generator,
 	globalStatus *status.GlobalStatus,
 	logger *zap.Logger,
-	stopFlag *stop.Flag,
 	failFast, _ bool,
 ) error {
 	schemaConfig := &schemaCfg
@@ -292,18 +273,20 @@ func warmupJob(
 		logger.Info("ending warmup loop")
 	}()
 	for {
-		if stopFlag.IsHardOrSoft() {
+		select {
+		case <-ctx.Done():
 			logger.Debug("warmup job terminated")
 			return nil
+		default:
 		}
+
 		// Do we care about errors during warmup?
 		err := mutation(ctx, schema, schemaConfig, table, s, r, p, g, globalStatus, false, logger)
 		if err != nil {
 			return err
 		}
 
 		if failFast && globalStatus.HasErrors() {
-			stopFlag.SetSoft(true)
 			return nil
 		}
 	}

diff --git a/pkg/jobs/pump.go b/pkg/jobs/pump.go
@@ -15,15 +15,14 @@
 package jobs
 
 import (
+	"context"
 	"time"
 
-	"github.com/scylladb/gemini/pkg/stop"
-
 	"go.uber.org/zap"
 	"golang.org/x/exp/rand"
 )
 
-func NewPump(stopFlag *stop.Flag, logger *zap.Logger) chan time.Duration {
+func NewPump(ctx context.Context, logger *zap.Logger) <-chan time.Duration {
 	pump := make(chan time.Duration, 10000)
 	logger = logger.Named("Pump")
 	go func() {
@@ -32,8 +31,14 @@ func NewPump(stopFlag *stop.Flag, logger *zap.Logger) chan time.Duration {
 			close(pump)
 			logger.Debug("pump channel closed")
 		}()
-		for !stopFlag.IsHardOrSoft() {
-			pump <- newHeartBeat()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				pump <- newHeartBeat()
+			}
 		}
 	}()