refactor: pass the correlationID direcly instead of using the context

2 files changed, 10 insertions, 12 deletions

diff --git a/internal/source/gitlab/cache/cache.go b/internal/source/gitlab/cache/cache.go
index 84c35a62..7dddb176 100644
--- a/internal/source/gitlab/cache/cache.go
+++ b/internal/source/gitlab/cache/cache.go
@@ -94,29 +94,24 @@ func (c *Cache) Resolve(ctx context.Context, domain string) *api.Lookup {
 	return c.retrieve(ctx, entry)
 }
 
-func (c *Cache) retrieve(originalCtx context.Context, entry *Entry) *api.Lookup {
+func (c *Cache) retrieve(ctx context.Context, entry *Entry) *api.Lookup {
 	// We run the code within an additional func() to run both `e.setResponse`
 	// and `c.retriever.Retrieve` asynchronously.
 	// We are using a sync.Once so this assumes that setResponse is always called
 	// the first (and only) time f is called, otherwise future requests will hang.
 	entry.retrieve.Do(func() {
-		// forward correlation_id from originalCtx to the new independent context
-		correlationID := correlation.ExtractFromContext(originalCtx)
-		retrieverCtx := correlation.ContextWithCorrelation(context.Background(), correlationID)
-
-		retrieverCtx, cancel := context.WithTimeout(retrieverCtx, c.retriever.retrievalTimeout)
+		correlationID := correlation.ExtractFromContext(ctx)
 
 		go func() {
-			l := c.retriever.Retrieve(retrieverCtx, entry.domain)
+			l := c.retriever.Retrieve(correlationID, entry.domain)
 			entry.setResponse(l)
-			cancel()
 		}()
 	})
 
 	var lookup *api.Lookup
 	select {
-	case <-originalCtx.Done():
-		lookup = &api.Lookup{Name: entry.domain, Error: fmt.Errorf("original context done: %w", originalCtx.Err())}
+	case <-ctx.Done():
+		lookup = &api.Lookup{Name: entry.domain, Error: fmt.Errorf("original context done: %w", ctx.Err())}
 	case <-entry.retrieved:
 		lookup = entry.Lookup()
 	}
diff --git a/internal/source/gitlab/cache/retriever.go b/internal/source/gitlab/cache/retriever.go
index d6dd0924..01ce4a16 100644
--- a/internal/source/gitlab/cache/retriever.go
+++ b/internal/source/gitlab/cache/retriever.go
@@ -35,10 +35,13 @@ func NewRetriever(client api.Client, retrievalTimeout, maxRetrievalInterval time
 
 // Retrieve retrieves a lookup response from external source with timeout and
 // backoff. It has its own context with timeout.
-func (r *Retriever) Retrieve(ctx context.Context, domain string) (lookup api.Lookup) {
+func (r *Retriever) Retrieve(correlationID, domain string) (lookup api.Lookup) {
 	var logMsg string
 
-	correlationID := correlation.ExtractFromContext(ctx)
+	ctx := correlation.ContextWithCorrelation(context.Background(), correlationID)
+
+	ctx, cancel := context.WithTimeout(ctx, r.retrievalTimeout)
+	defer cancel()
 
 	select {
 	case <-ctx.Done():