Really really make sure SAs get added (#108)

* Really really make sure SAs get added

* Deal with unit test

cmd/run_scenario.go

@@ -3,8 +3,6 @@ package main
 import (
 	"context"
 	"fmt"
-	"go.temporal.io/api/enums/v1"
-	"go.temporal.io/api/operatorservice/v1"
 	"strings"
 	"time"
 
@@ -124,18 +122,6 @@ func (r *scenarioRunner) run(ctx context.Context) error {
 	}
 	defer client.Close()
 
-	// Ensure required search attributes are registered
-	_, err = client.OperatorService().AddSearchAttributes(ctx, &operatorservice.AddSearchAttributesRequest{
-		SearchAttributes: map[string]enums.IndexedValueType{
-			"KS_Keyword": enums.INDEXED_VALUE_TYPE_KEYWORD,
-			"KS_Int":     enums.INDEXED_VALUE_TYPE_INT,
-		},
-		Namespace: r.clientOptions.Namespace,
-	})
-	if err != nil {
-		r.logger.Warnf("Error ignored when registering search attributes: %v", err)
-	}
-
 	scenarioInfo := loadgen.ScenarioInfo{
 		ScenarioName:   r.scenario,
 		RunID:          r.runID,

loadgen/generic_executor.go

@@ -71,6 +71,13 @@ func (g *GenericExecutor) newRun(info ScenarioInfo) (*genericRun, error) {
 // iterations is reached.
 func (g *genericRun) Run(ctx context.Context) error {
 	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	err := g.info.RegisterDefaultSearchAttributes(ctx)
+	if err != nil {
+		return err
+	}
+
 	if g.config.Timeout > 0 {
 		g.logger.Debugf("Will timeout after %v", g.config.Timeout)
 		ctx, cancel = context.WithTimeout(ctx, g.config.Timeout)

loadgen/kitchensink/helpers.go

@@ -80,7 +80,7 @@ func (e *ClientActionsExecutor) executeClientActionSet(ctx context.Context, acti
 			errs.Go(func() error {
 				err := e.executeClientAction(errGroupCtx, action)
 				if err != nil {
-					return fmt.Errorf("failed to execute client action %v: %w", action, err)
+					return fmt.Errorf("failed to execute concurrent client action %v: %w", action, err)
 				}
 				return nil
 			})

loadgen/scenario.go

@@ -3,6 +3,8 @@ package loadgen
 import (
 	"context"
 	"fmt"
+	"go.temporal.io/api/enums/v1"
+	"go.temporal.io/api/operatorservice/v1"
 	"path/filepath"
 	"runtime"
 	"strconv"
@@ -157,6 +159,40 @@ func (s *ScenarioInfo) NewRun(iteration int) *Run {
 	}
 }
 
+func (s *ScenarioInfo) RegisterDefaultSearchAttributes(ctx context.Context) error {
+	if s.Client == nil {
+		// No client in some unit tests. Ideally this would be mocked but no mock operator service
+		// client is readily available.
+		return nil
+	}
+	// Ensure custom search attributes are registered that many scenarios rely on
+	_, err := s.Client.OperatorService().AddSearchAttributes(ctx, &operatorservice.AddSearchAttributesRequest{
+		SearchAttributes: map[string]enums.IndexedValueType{
+			"KS_Keyword": enums.INDEXED_VALUE_TYPE_KEYWORD,
+			"KS_Int":     enums.INDEXED_VALUE_TYPE_INT,
+		},
+		Namespace: s.Namespace,
+	})
+	// Throw an error if the attributes could not be registered, but ignore already exists errs
+	alreadyExistsStrings := []string{
+		"already exists",
+		"attributes mapping unavailble",
+	}
+	if err != nil {
+		isAlreadyExistsErr := false
+		for _, s := range alreadyExistsStrings {
+			if strings.Contains(err.Error(), s) {
+				isAlreadyExistsErr = true
+				break
+			}
+		}
+		if !isAlreadyExistsErr {
+			return fmt.Errorf("failed to register search attributes: %w", err)
+		}
+	}
+	return nil
+}
+
 // TaskQueueForRun returns a default task queue name for the given scenario name and run ID.
 func TaskQueueForRun(scenarioName, runID string) string {
 	return fmt.Sprintf("%s:%s", scenarioName, runID)