From d114f3646ddf948935f33e50b27f62728f22b7f8 Mon Sep 17 00:00:00 2001
From: Earl Warren <contact@earl-warren.org>
Date: Sat, 16 Aug 2025 19:13:32 +0200
Subject: [PATCH] feat: insert the daemon context in the poller context

The daemon context is needed when the context of a job or the poller
is done. Otherwise it is no longer possible to send a conclusion
report to Forgejo, short of creating a context.Background() which
poses its own set of problems.

- WithDaemonContext is used to store the daemon context
- The poller uses the daemon context instead of context.Background
---
 internal/app/cmd/cmd.go               |  3 ++-
 internal/app/cmd/daemon.go            |  2 +-
 internal/app/poll/poller.go           | 10 +++++-----
 internal/app/poll/poller_test.go      |  4 +++-
 internal/pkg/common/daemon_context.go | 26 ++++++++++++++++++++++++++
 5 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 internal/pkg/common/daemon_context.go

diff --git a/internal/app/cmd/cmd.go b/internal/app/cmd/cmd.go
index c3a98a88..2f9bd3e4 100644
--- a/internal/app/cmd/cmd.go
+++ b/internal/app/cmd/cmd.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"code.forgejo.org/forgejo/runner/v9/internal/pkg/common"
 	"code.forgejo.org/forgejo/runner/v9/internal/pkg/config"
 	"code.forgejo.org/forgejo/runner/v9/internal/pkg/ver"
 )
@@ -45,7 +46,7 @@ func Execute(ctx context.Context) {
 		Use:   "daemon",
 		Short: "Run as a runner daemon",
 		Args:  cobra.MaximumNArgs(1),
-		RunE:  runDaemon(ctx, &configFile),
+		RunE:  runDaemon(common.WithDaemonContext(ctx, ctx), &configFile),
 	}
 	rootCmd.AddCommand(daemonCmd)
 
diff --git a/internal/app/cmd/daemon.go b/internal/app/cmd/daemon.go
index 0a9815ed..22fd8c26 100644
--- a/internal/app/cmd/daemon.go
+++ b/internal/app/cmd/daemon.go
@@ -114,7 +114,7 @@ func runDaemon(ctx context.Context, configFile *string) func(cmd *cobra.Command,
 			}
 		}
 
-		poller := poll.New(cfg, cli, runner)
+		poller := poll.New(ctx, cfg, cli, runner)
 
 		go poller.Poll()
 
diff --git a/internal/app/poll/poller.go b/internal/app/poll/poller.go
index 4ebe4405..5a4061f8 100644
--- a/internal/app/poll/poller.go
+++ b/internal/app/poll/poller.go
@@ -42,14 +42,14 @@ type poller struct {
 	done chan any
 }
 
-func New(cfg *config.Config, client client.Client, runner run.RunnerInterface) Poller {
-	return (&poller{}).init(cfg, client, runner)
+func New(ctx context.Context, cfg *config.Config, client client.Client, runner run.RunnerInterface) Poller {
+	return (&poller{}).init(ctx, cfg, client, runner)
 }
 
-func (p *poller) init(cfg *config.Config, client client.Client, runner run.RunnerInterface) Poller {
-	pollingCtx, shutdownPolling := context.WithCancel(context.Background())
+func (p *poller) init(ctx context.Context, cfg *config.Config, client client.Client, runner run.RunnerInterface) Poller {
+	pollingCtx, shutdownPolling := context.WithCancel(ctx)
 
-	jobsCtx, shutdownJobs := context.WithCancel(context.Background())
+	jobsCtx, shutdownJobs := context.WithCancel(ctx)
 
 	done := make(chan any)
 
diff --git a/internal/app/poll/poller_test.go b/internal/app/poll/poller_test.go
index 0c5febfe..7c8b7633 100644
--- a/internal/app/poll/poller_test.go
+++ b/internal/app/poll/poller_test.go
@@ -118,7 +118,7 @@ func setTrace(t *testing.T) {
 }
 
 func TestPoller_New(t *testing.T) {
-	p := New(&config.Config{}, &mockClient{}, &mockRunner{})
+	p := New(t.Context(), &config.Config{}, &mockClient{}, &mockRunner{})
 	assert.NotNil(t, p)
 }
 
@@ -172,6 +172,7 @@ func TestPoller_Runner(t *testing.T) {
 			}
 			p := &mockPoller{}
 			p.init(
+				t.Context(),
 				&config.Config{
 					Runner: configRunner,
 				},
@@ -239,6 +240,7 @@ func TestPoller_Fetch(t *testing.T) {
 			}
 			p := &mockPoller{}
 			p.init(
+				t.Context(),
 				&config.Config{
 					Runner: configRunner,
 				},
diff --git a/internal/pkg/common/daemon_context.go b/internal/pkg/common/daemon_context.go
new file mode 100644
index 00000000..8cb2f70a
--- /dev/null
+++ b/internal/pkg/common/daemon_context.go
@@ -0,0 +1,26 @@
+// Copyright 2025 The Forgejo Authors
+// SPDX-License-Identifier: MIT
+
+package common
+
+import (
+	"context"
+)
+
+type daemonContextKey string
+
+const daemonContextKeyVal = daemonContextKey("daemon")
+
+func DaemonContext(ctx context.Context) context.Context {
+	val := ctx.Value(daemonContextKeyVal)
+	if val != nil {
+		if daemon, ok := val.(context.Context); ok {
+			return daemon
+		}
+	}
+	return context.Background()
+}
+
+func WithDaemonContext(ctx, daemon context.Context) context.Context {
+	return context.WithValue(ctx, daemonContextKeyVal, daemon)
+}