gemini: added CQL tracing support

A new CLI argument "tracing-outfile" is added that denotes a file
where tracing information is saved. By default tracing is disabled.

cmd/gemini/root.go

@@ -64,6 +64,7 @@ var (
 	partitionKeyDistribution string
 	normalDistMean           float64
 	normalDistSigma          float64
+	tracingOutFile           string
 )
 
 const (
@@ -136,19 +137,14 @@ func run(cmd *cobra.Command, args []string) error {
 	if err := printSetup(); err != nil {
 		return errors.Wrapf(err, "unable to print setup")
 	}
-
 	distFunc, err := createDistributionFunc(partitionKeyDistribution, distributionSize, seed, stdDistMean, oneStdDev)
 	if err != nil {
 		return err
 	}
 
-	outFile := os.Stdout
-	if outFileArg != "" {
-		of, err := os.Create(outFileArg)
-		if err != nil {
-			return errors.Wrapf(err, "Unable to open output file %s", outFileArg)
-		}
-		outFile = of
+	outFile, err := createFile(outFileArg, os.Stdout)
+	if err != nil {
+		return err
 	}
 	defer outFile.Sync()
 
@@ -175,7 +171,23 @@ func run(cmd *cobra.Command, args []string) error {
 		MaxRetriesMutate:      maxRetriesMutate,
 		MaxRetriesMutateSleep: maxRetriesMutateSleep,
 	}
-	store := store.New(schema, testCluster, oracleCluster, storeConfig, logger)
+	var tracingFile *os.File
+	if tracingOutFile != "" {
+		switch tracingOutFile {
+		case "stderr":
+			tracingFile = os.Stderr
+		case "stdout":
+			tracingFile = os.Stdout
+		default:
+			tf, err := createFile(tracingOutFile, os.Stdout)
+			if err != nil {
+				return err
+			}
+			tracingFile = tf
+			defer tracingFile.Sync()
+		}
+	}
+	store := store.New(schema, testCluster, oracleCluster, storeConfig, tracingFile, logger)
 	defer store.Close()
 
 	if dropSchema && mode != readMode {
@@ -220,6 +232,17 @@ func run(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+func createFile(fname string, def *os.File) (*os.File, error) {
+	if fname != "" {
+		f, err := os.Create(fname)
+		if err != nil {
+			return nil, errors.Wrapf(err, "Unable to open output file %s", fname)
+		}
+		return f, nil
+	}
+	return def, nil
+}
+
 const (
 	stdDistMean = 0.5
 	oneStdDev   = 0.341
@@ -400,6 +423,7 @@ func init() {
 	rootCmd.Flags().StringVarP(&partitionKeyDistribution, "partition-key-distribution", "", "uniform", "Specify the distribution from which to draw partition keys, supported values are currently uniform|normal|exponential")
 	rootCmd.Flags().Float64VarP(&normalDistMean, "normal-dist-mean", "", stdDistMean, "Mean of the normal distribution")
 	rootCmd.Flags().Float64VarP(&normalDistSigma, "normal-dist-sigma", "", oneStdDev, "Sigma of the normal distribution, defaults to one standard deviation ~0.341")
+	rootCmd.Flags().StringVarP(&tracingOutFile, "tracing-outfile", "", "", "Specify the file to which tracing information gets written. Two magic names are available, 'stdout' and 'stderr'. By default tracing is disabled.")
 }
 
 func printSetup() error {

store/cqlstore.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"context"
+	"io"
 	"time"
 
 	"github.com/gocql/gocql"
@@ -100,11 +101,14 @@ func (cs cqlStore) close() error {
 	return nil
 }
 
-func newSession(cluster *gocql.ClusterConfig) *gocql.Session {
+func newSession(cluster *gocql.ClusterConfig, out io.Writer) *gocql.Session {
 	session, err := cluster.CreateSession()
 	if err != nil {
 		panic(err)
 	}
+	if out != nil {
+		session.SetTrace(gocql.NewTraceWriter(session, out))
+	}
 	return session
 }
 

store/store.go

@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"fmt"
+	"io"
 	"math/big"
 	"sort"
 	"time"
@@ -48,7 +49,7 @@ type Config struct {
 	MaxRetriesMutateSleep time.Duration
 }
 
-func New(schema *gemini.Schema, testCluster *gocql.ClusterConfig, oracleCluster *gocql.ClusterConfig, cfg Config, logger *zap.Logger) Store {
+func New(schema *gemini.Schema, testCluster *gocql.ClusterConfig, oracleCluster *gocql.ClusterConfig, cfg Config, traceOut io.Writer, logger *zap.Logger) Store {
 	ops := promauto.NewCounterVec(prometheus.CounterOpts{
 		Name: "gemini_cql_requests",
 		Help: "How many CQL requests processed, partitioned by system and CQL query type aka 'method' (batch, delete, insert, update).",
@@ -59,7 +60,7 @@ func New(schema *gemini.Schema, testCluster *gocql.ClusterConfig, oracleCluster
 	var validations bool
 	if oracleCluster != nil {
 		oracleStore = &cqlStore{
-			session:               newSession(oracleCluster),
+			session:               newSession(oracleCluster, traceOut),
 			schema:                schema,
 			system:                "oracle",
 			ops:                   ops,
@@ -76,7 +77,7 @@ func New(schema *gemini.Schema, testCluster *gocql.ClusterConfig, oracleCluster
 
 	return &delegatingStore{
 		testStore: &cqlStore{
-			session:               newSession(testCluster),
+			session:               newSession(testCluster, traceOut),
 			schema:                schema,
 			system:                "test",
 			ops:                   ops,