Merge pull request #1869 from pikasTech/fix-v03-workbench-runtime-query-timeout-1868

fix(workbench-runtime): bound sessions summary query retries

internal/workbenchruntime/service.go

@@ -20,6 +20,12 @@ import (
 
 const serviceID = "hwlab-workbench-runtime"
 
+const (
+	sessionSummaryQueryMaxAttempts    = 2
+	sessionSummaryQueryAttemptTimeout = 2 * time.Second
+	sessionSummaryQueryRetryBackoff   = 150 * time.Millisecond
+)
+
 type Config struct {
 	Addr           string
 	DatabaseURL    string
@@ -545,8 +551,41 @@ func (s *Server) dbQueryAttrs(table string, argCount int, extra map[string]any)
 }
 
 func (s *Server) querySessionSummaries(ctx context.Context, stage string, sqlText string, args []any) ([]any, error) {
+	var lastErr error
+	for attempt := 1; attempt <= sessionSummaryQueryMaxAttempts; attempt++ {
+		attemptResult, err := s.querySessionSummariesAttempt(ctx, stage, sqlText, args, attempt, sessionSummaryQueryMaxAttempts)
+		if err == nil {
+			return attemptResult, nil
+		}
+		lastErr = err
+		if attempt == sessionSummaryQueryMaxAttempts || !retryableSessionSummaryQueryError(err) || ctx.Err() != nil {
+			break
+		}
+		backoff := sessionSummaryRetryBackoff(attempt)
+		timer := time.NewTimer(backoff)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return nil, wrapQueryStage(stage, ctx.Err())
+		case <-timer.C:
+		}
+	}
+	return nil, wrapQueryStage(stage, lastErr)
+}
+
+func (s *Server) querySessionSummariesAttempt(ctx context.Context, stage string, sqlText string, args []any, attempt int, maxAttempts int) ([]any, error) {
 	result := []any{}
-	err := s.withOtelInternalSpan(ctx, stage, s.dbQueryAttrs("workbench_sessions", len(args), map[string]any{"db.index.expected": "idx_workbench_sessions_owner_updated|idx_workbench_sessions_updated"}), func(spanCtx context.Context) error {
+	attemptCtx, cancel := context.WithTimeout(ctx, sessionSummaryQueryAttemptTimeout)
+	defer cancel()
+	attrs := s.dbQueryAttrs("workbench_sessions", len(args), map[string]any{
+		"db.index.expected": "idx_workbench_sessions_owner_updated|idx_workbench_sessions_updated",
+		"retryAttempt":      int64(attempt),
+		"retryMax":          int64(maxAttempts),
+	})
+	if attempt > 1 {
+		attrs["retryBackoffMs"] = sessionSummaryRetryBackoff(attempt - 1).Milliseconds()
+	}
+	err := s.withOtelInternalSpan(attemptCtx, stage, attrs, func(spanCtx context.Context) error {
 		rows, err := s.db.QueryContext(spanCtx, sqlText, args...)
 		if err != nil {
 			return err
@@ -579,10 +618,29 @@ func (s *Server) querySessionSummaries(ctx context.Context, stage string, sqlTex
 		}
 		return rows.Err()
 	})
-	if err != nil {
+	return result, err
-		return nil, wrapQueryStage(stage, err)
+}
+
+func sessionSummaryRetryBackoff(attempt int) time.Duration {
+	if attempt <= 1 {
+		return sessionSummaryQueryRetryBackoff
+	}
+	return time.Duration(attempt) * sessionSummaryQueryRetryBackoff
+}
+
+func retryableSessionSummaryQueryError(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+		return true
+	}
+	switch sqlErrorCode(err) {
+	case "40001", "40P01", "55P03", "57014", "query_timeout":
+		return true
+	default:
+		return false
 	}
-	return result, nil
 }
 
 func (s *Server) queryMessageSummaries(ctx context.Context, sessionIDs []string) ([]any, error) {