scheduler/state: avoid passing JobStore via ctx

This (hack) is no longer necessary since all indexing logic is now self-contained within one package.
hashicorp · Jul 5, 2022 · 8ef3b35 · 8ef3b35
1 parent 0b351b2
commit 8ef3b35
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 27 deletions.
diff --git a/internal/job/job_store.go b/internal/job/job_store.go
@@ -2,24 +2,9 @@ package job
 
 import (
 	"context"
-	"fmt"
 )
 
 type JobStore interface {
 	EnqueueJob(newJob Job) (ID, error)
 	WaitForJobs(ctx context.Context, ids ...ID) error
 }
-
-type jobStoreCtxKey struct{}
-
-func WithJobStore(ctx context.Context, js JobStore) context.Context {
-	return context.WithValue(ctx, jobStoreCtxKey{}, js)
-}
-
-func JobStoreFromContext(ctx context.Context) (JobStore, error) {
-	js, ok := ctx.Value(jobStoreCtxKey{}).(JobStore)
-	if !ok {
-		return nil, fmt.Errorf("not found JobStore in context")
-	}
-	return js, nil
-}
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
@@ -69,8 +69,7 @@ func (s *Scheduler) eval(ctx context.Context) {
 
 		deferredJobIds := make(job.IDs, 0)
 		if nextJob.Defer != nil {
-			deferCtx := job.WithJobStore(ctx, s.jobStorage)
-			deferredJobIds, err = nextJob.Defer(deferCtx, jobErr)
+			deferredJobIds, err = nextJob.Defer(ctx, jobErr)
 			if err != nil {
 				s.logger.Printf("deferred job failed: %s", err)
 			}

diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -262,11 +262,7 @@ func TestScheduler_defer(t *testing.T) {
 			Type: "test-type",
 			Defer: func(ctx context.Context, jobErr error) (job.IDs, error) {
 				ids := make(job.IDs, 0)
-				je, err := job.JobStoreFromContext(ctx)
-				if err != nil {
-					log.Fatal(err)
-					return ids, err
-				}
+				je := ss.JobStore
 
 				id1, err := je.EnqueueJob(job.Job{
 					Dir:  document.DirHandleFromPath(dirPath),

diff --git a/internal/state/jobs_test.go b/internal/state/jobs_test.go
@@ -640,10 +640,7 @@ func TestJobStore_FinishJob_defer(t *testing.T) {
 
 	defer1Func := func(ctx context.Context, jobErr error) (job.IDs, error) {
 		ids := make(job.IDs, 0)
-		jobStore, err := job.JobStoreFromContext(ctx)
-		if err != nil {
-			return ids, err
-		}
+		jobStore := ss.JobStore
 
 		id, err := jobStore.EnqueueJob(job.Job{
 			Func: func(ctx context.Context) error {
@@ -672,7 +669,6 @@ func TestJobStore_FinishJob_defer(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = job.WithJobStore(ctx, ss.JobStore)
 	// execute deferred func, which is what scheduler would do
 	deferredIds, err := defer1Func(ctx, nil)
 	if err != nil {