chore(cli): scaletest: move logic for flushing traces into tracing provider

johnstcn · johnstcn · commit b105e67dec90 · 2023-05-04T13:42:24.000+01:00
diff --git a/cli/scaletest.go b/cli/scaletest.go
@@ -109,7 +109,10 @@ func (s *scaletestTracingFlags) provider(ctx context.Context) (trace.TracerProvi
 	return tracerProvider, func(ctx context.Context) error {
 		var err error
 		closeTracingOnce.Do(func() {
-			err = closeTracing(ctx)
+			// Allow time to upload traces even if ctx is canceled
+			traceCtx, traceCancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer traceCancel()
+			err = closeTracing(traceCtx)
 		})
 
 		return err
@@ -637,10 +640,11 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 			}
 			defer func() {
 				// Allow time for traces to flush even if command context is
-				// canceled.
-				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-				defer cancel()
-				_ = closeTracing(ctx)
+				// canceled. This is a no-op if tracing is not enabled.
+				_, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...")
+				if err := closeTracing(ctx); err != nil {
+					_, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err)
+				}
 			}()
 			tracer := tracerProvider.Tracer(scaletestTracerName)
 
@@ -754,17 +758,6 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 				return xerrors.Errorf("cleanup tests: %w", err)
 			}
 
-			// Upload traces.
-			if tracingEnabled {
-				_, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...")
-				ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
-				defer cancel()
-				err := closeTracing(ctx)
-				if err != nil {
-					_, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err)
-				}
-			}
-
 			if res.TotalFail > 0 {
 				return xerrors.New("load test failed, see above for more details")
 			}
@@ -905,7 +898,7 @@ func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 	var (
 		duration        time.Duration
 		bps             int64
-		client          = new(codersdk.Client)
+		client          = &codersdk.Client{}
 		tracingFlags    = &scaletestTracingFlags{}
 		strategy        = &scaletestStrategyFlags{}
 		cleanupStrategy = &scaletestStrategyFlags{cleanup: true}
@@ -947,10 +940,11 @@ func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 			}
 			defer func() {
 				// Allow time for traces to flush even if command context is
-				// canceled.
-				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-				defer cancel()
-				_ = closeTracing(ctx)
+				// canceled. This is a no-op if tracing is not enabled.
+				_, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...")
+				if err := closeTracing(ctx); err != nil {
+					_, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err)
+				}
 			}()
 			tracer := tracerProvider.Tracer(scaletestTracerName)
 
@@ -1018,17 +1012,6 @@ func (r *RootCmd) scaletestTrafficGen() *clibase.Cmd {
 				}
 			}
 
-			// Upload traces.
-			if tracingEnabled {
-				_, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...")
-				ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
-				defer cancel()
-				err := closeTracing(ctx)
-				if err != nil {
-					_, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err)
-				}
-			}
-
 			if res.TotalFail > 0 {
 				return xerrors.New("load test failed, see above for more details")
 			}
