feat: add granular context timeout (#112952)

* feat: add granular context timeout

* test: forced thread timeout

* fix: test assumption for blocks

* test: new change

* trigger build

* remove loggers to test

* test with fmt.print

* fix: deadlocks

pkg/registry/apis/provisioning/jobs/progress.go

@@ -70,23 +70,31 @@ func newJobProgressRecorder(ProgressFn ProgressFn) JobProgressRecorder {
 }
 
 func (r *jobProgressRecorder) Record(ctx context.Context, result JobResourceResult) {
+	var shouldLogError bool
+	var logErr error
+
 	r.mu.Lock()
 	r.resultCount++
 
-	logger := logging.FromContext(ctx).With("path", result.Path, "group", result.Group, "kind", result.Kind, "action", result.Action, "name", result.Name)
 	if result.Error != nil {
-		logger.Error("job resource operation failed", "err", result.Error)
+		shouldLogError = true
+		logErr = result.Error
 		if len(r.errors) < 20 {
 			r.errors = append(r.errors, result.Error.Error())
 		}
 		r.errorCount++
-	} else {
-		logger.Info("job resource operation succeeded")
 	}
 
 	r.updateSummary(result)
 	r.mu.Unlock()
 
+	logger := logging.FromContext(ctx).With("path", result.Path, "group", result.Group, "kind", result.Kind, "action", result.Action, "name", result.Name)
+	if shouldLogError {
+		logger.Error("job resource operation failed", "err", logErr)
+	} else {
+		logger.Info("job resource operation succeeded")
+	}
+
 	r.maybeNotify(ctx)
 }
 
@@ -145,9 +153,6 @@ func (r *jobProgressRecorder) StrictMaxErrors(maxErrors int) {
 }
 
 func (r *jobProgressRecorder) TooManyErrors() error {
-	r.mu.RLock()
-	defer r.mu.RUnlock()
-
 	if r.maxErrors > 0 && r.errorCount >= r.maxErrors {
 		return fmt.Errorf("too many errors: %d", r.errorCount)
 	}
@@ -156,9 +161,6 @@ func (r *jobProgressRecorder) TooManyErrors() error {
 }
 
 func (r *jobProgressRecorder) summary() []*provisioning.JobResourceSummary {
-	r.mu.RLock()
-	defer r.mu.RUnlock()
-
 	if len(r.summaries) == 0 {
 		return nil
 	}
@@ -247,13 +249,9 @@ func (r *jobProgressRecorder) maybeNotify(ctx context.Context) {
 
 func (r *jobProgressRecorder) Complete(ctx context.Context, err error) provisioning.JobStatus {
 	r.mu.RLock()
-	defer r.mu.RUnlock()
 
-	// Initialize base job status
 	jobStatus := provisioning.JobStatus{
-		Started: r.started.UnixMilli(),
-		// FIXME: if we call this method twice, the state will be different
-		// This results in sync status to be different from job status
+		Started:  r.started.UnixMilli(),
 		Finished: time.Now().UnixMilli(),
 		State:    provisioning.JobStateSuccess,
 		Message:  "completed successfully",
@@ -268,9 +266,13 @@ func (r *jobProgressRecorder) Complete(ctx context.Context, err error) provision
 	jobStatus.Errors = r.errors
 	jobStatus.URLs = r.refURLs
 
-	// Check for errors during execution
+	tooManyErrors := r.maxErrors > 0 && r.errorCount >= r.maxErrors
+	finalMessage := r.finalMessage
+
+	r.mu.RUnlock()
+
 	if len(jobStatus.Errors) > 0 && jobStatus.State != provisioning.JobStateError {
-		if r.TooManyErrors() != nil {
+		if tooManyErrors {
 			jobStatus.Message = "completed with too many errors"
 			jobStatus.State = provisioning.JobStateError
 		} else {
@@ -280,8 +282,8 @@ func (r *jobProgressRecorder) Complete(ctx context.Context, err error) provision
 	}
 
 	// Override message if progress have a more explicit message
-	if r.finalMessage != "" && jobStatus.State != provisioning.JobStateError {
-		jobStatus.Message = r.finalMessage
+	if finalMessage != "" && jobStatus.State != provisioning.JobStateError {
+		jobStatus.Message = finalMessage
 	}
 
 	return jobStatus

pkg/registry/apis/provisioning/jobs/sync/full.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 
+	"github.com/grafana/grafana-app-sdk/logging"
 	"github.com/grafana/grafana/apps/provisioning/pkg/repository"
 	"github.com/grafana/grafana/apps/provisioning/pkg/safepath"
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -145,11 +147,11 @@ func applyChange(ctx context.Context, change ResourceFileChange, clients resourc
 		Group:  gvk.Group,
 		Kind:   gvk.Kind,
 	}
-
 	if err != nil {
 		writeSpan.RecordError(err)
 		result.Error = fmt.Errorf("writing resource from file %s: %w", change.Path, err)
 	}
+
 	progress.Record(writeCtx, result)
 	writeSpan.End()
 }
@@ -224,72 +226,75 @@ func applyChanges(ctx context.Context, changes []ResourceFileChange, clients res
 }
 
 func applyFoldersSerially(ctx context.Context, folders []ResourceFileChange, clients resources.ResourceClients, repositoryResources resources.RepositoryResources, progress jobs.JobProgressRecorder, tracer tracing.Tracer) error {
-	folderCtx, folderCancel := context.WithCancel(ctx)
-	defer folderCancel()
+	logger := logging.FromContext(ctx)
 
 	for _, folder := range folders {
-		if folderCtx.Err() != nil {
-			return folderCtx.Err()
+		if ctx.Err() != nil {
+			return ctx.Err()
 		}
 
 		if err := progress.TooManyErrors(); err != nil {
 			return err
 		}
 
+		folderCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
+
 		applyChange(folderCtx, folder, clients, repositoryResources, progress, tracer)
+
+		if folderCtx.Err() == context.DeadlineExceeded {
+			logger.Error("operation timed out after 15 seconds", "path", folder.Path, "action", folder.Action)
+
+			recordCtx, recordCancel := context.WithTimeout(context.Background(), 15*time.Second)
+			progress.Record(recordCtx, jobs.JobResourceResult{
+				Path:   folder.Path,
+				Action: folder.Action,
+				Error:  fmt.Errorf("operation timed out after 15 seconds"),
+			})
+			recordCancel()
+		}
+
+		cancel()
 	}
 
 	return nil
 }
 
 func applyResourcesInParallel(ctx context.Context, resources []ResourceFileChange, clients resources.ResourceClients, repositoryResources resources.RepositoryResources, progress jobs.JobProgressRecorder, tracer tracing.Tracer, maxSyncWorkers int) error {
+	logger := logging.FromContext(ctx)
+	logger.Info("applying resources in parallel test changes 1")
+
 	if len(resources) == 0 {
 		return nil
 	}
 
-	workerCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	changeChan := make(chan ResourceFileChange, len(resources))
+	sem := make(chan struct{}, maxSyncWorkers)
 	var wg sync.WaitGroup
 
-	for i := 0; i < maxSyncWorkers; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			for {
-				select {
-				case change, ok := <-changeChan:
-					if !ok {
-						return
-					}
-
-					if err := progress.TooManyErrors(); err != nil {
-						cancel()
-						return
-					}
-					if workerCtx.Err() != nil {
-						return
-					}
-
-					applyChange(workerCtx, change, clients, repositoryResources, progress, tracer)
-
-				case <-workerCtx.Done():
-					return
-				}
-			}
-		}()
-	}
-
+loop:
 	for _, change := range resources {
-		select {
-		case changeChan <- change:
-		case <-workerCtx.Done():
-			goto done
+		if err := progress.TooManyErrors(); err != nil {
+			break
 		}
+		if ctx.Err() != nil {
+			break
+		}
+
+		// Acquire semaphore slot (blocks if max workers reached)
+		select {
+		case sem <- struct{}{}:
+		case <-ctx.Done():
+			break loop
+		}
+
+		wg.Add(1)
+		go func(change ResourceFileChange) {
+			defer wg.Done()
+			defer func() { <-sem }()
+
+			applyChangeWithTimeout(ctx, change, clients, repositoryResources, progress, tracer, logger)
+		}(change)
 	}
-done:
-	close(changeChan)
+
 	wg.Wait()
 
 	if err := progress.TooManyErrors(); err != nil {
@@ -298,3 +303,22 @@ done:
 
 	return ctx.Err()
 }
+
+func applyChangeWithTimeout(ctx context.Context, change ResourceFileChange, clients resources.ResourceClients, repositoryResources resources.RepositoryResources, progress jobs.JobProgressRecorder, tracer tracing.Tracer, logger logging.Logger) {
+	changeCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
+	defer cancel()
+
+	applyChange(changeCtx, change, clients, repositoryResources, progress, tracer)
+
+	if changeCtx.Err() == context.DeadlineExceeded {
+		logger.Error("operation timed out after 15 seconds", "path", change.Path, "action", change.Action)
+
+		recordCtx, recordCancel := context.WithTimeout(context.Background(), 15*time.Second)
+		progress.Record(recordCtx, jobs.JobResourceResult{
+			Path:   change.Path,
+			Action: change.Action,
+			Error:  fmt.Errorf("operation timed out after 15 seconds"),
+		})
+		recordCancel()
+	}
+}

pkg/registry/apis/provisioning/jobs/sync/full_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -202,10 +204,9 @@ func TestFullSync_ApplyChanges(t *testing.T) { //nolint:gocyclo
 				},
 			},
 			setupMocks: func(repo *repository.MockRepository, repoResources *resources.MockRepositoryResources, clients *resources.MockResourceClients, progress *jobs.MockJobProgressRecorder, compareFn *MockCompareFn) {
-				callCount := 0
+				var callCount int64 = 0
 				progress.On("TooManyErrors").Return(func() error {
-					callCount++
-					if callCount > 1 {
+					if atomic.AddInt64(&callCount, 1) > 1 {
 						return fmt.Errorf("too many errors")
 					}
 					return nil
@@ -682,6 +683,45 @@ func TestFullSync_ApplyChanges(t *testing.T) { //nolint:gocyclo
 				})).Return()
 			},
 		},
+		{
+			name:        "operation timeout after 15 seconds",
+			description: "Should record timeout error when operation takes longer than 15 seconds",
+			changes: []ResourceFileChange{
+				{
+					Action: repository.FileActionCreated,
+					Path:   "dashboards/slow.json",
+				},
+			},
+			setupMocks: func(repo *repository.MockRepository, repoResources *resources.MockRepositoryResources, clients *resources.MockResourceClients, progress *jobs.MockJobProgressRecorder, compareFn *MockCompareFn) {
+				progress.On("TooManyErrors").Return(nil)
+
+				repoResources.On("WriteResourceFromFile", mock.Anything, "dashboards/slow.json", "").
+					Run(func(args mock.Arguments) {
+						ctx := args.Get(0).(context.Context)
+						select {
+						case <-ctx.Done():
+							return
+						case <-time.After(20 * time.Second):
+							return
+						}
+					}).
+					Return("", schema.GroupVersionKind{}, context.DeadlineExceeded)
+
+				progress.On("Record", mock.Anything, mock.MatchedBy(func(result jobs.JobResourceResult) bool {
+					return result.Action == repository.FileActionCreated &&
+						result.Path == "dashboards/slow.json" &&
+						result.Error != nil &&
+						result.Error.Error() == "writing resource from file dashboards/slow.json: context deadline exceeded"
+				})).Return().Once()
+
+				progress.On("Record", mock.Anything, mock.MatchedBy(func(result jobs.JobResourceResult) bool {
+					return result.Action == repository.FileActionCreated &&
+						result.Path == "dashboards/slow.json" &&
+						result.Error != nil &&
+						result.Error.Error() == "operation timed out after 15 seconds"
+				})).Return().Once()
+			},
+		},
 	}
 
 	for _, tt := range tests {