fix: clear consumed ephemeral registration state (#1256)

Ephemeral registrations are single-use, but the daemon kept `.runner` after the task, so a restarted container failed `Declare` with `unregistered runner` forever.

An ephemeral runner now removes its registration file as soon as it fetches a task, so a container killed mid-job registers afresh on restart. It also removes the file when Gitea answers `unregistered runner` at startup or while polling. A proxy's 401 does not count. Persistent runners keep their file, because Gitea sends the same reply for a token mismatch.

Closes https://gitea.com/gitea/runner/issues/818

---------

Co-authored-by: silverwind <me@silverwind.io>
Reviewed-on: https://gitea.com/gitea/runner/pulls/1256
Reviewed-by: silverwind <2021+silverwind@noreply.gitea.com>

e2e/ephemeral_test.go

@@ -22,7 +22,7 @@ func testEphemeralRunner(t *testing.T) {
 	case <-time.After(5 * time.Second):
 		t.Fatal("ephemeral runner was not deleted")
 	}
-	if !poller.Unregistered() {
+	if !poller.RegistrationRejected() {
 		t.Fatal("ephemeral runner stopped without being deleted")
 	}
 }

internal/app/cmd/daemon.go

@@ -191,6 +191,9 @@ func runDaemon(ctx context.Context, daemArgs *daemonArgs, configFile *string) fu
 			return err
 		case err != nil:
 			log.WithError(err).Error("fail to invoke Declare")
+			if client.IsRegistrationRejected(err) {
+				removeEphemeralRegistration(reg, cfg.Runner.File)
+			}
 			return err
 		default:
 			log.Infof("runner: %s, with version: %s, with labels: %v, declare successfully",
@@ -215,7 +218,7 @@ func runDaemon(ctx context.Context, daemArgs *daemonArgs, configFile *string) fu
 			done := make(chan struct{})
 			go func() {
 				defer close(done)
-				poller.PollOnce()
+				poller.PollOnce(func() { removeEphemeralRegistration(reg, cfg.Runner.File) })
 			}()
 
 			// shutdown when we complete a job or cancel is requested
@@ -244,6 +247,10 @@ func runDaemon(ctx context.Context, daemArgs *daemonArgs, configFile *string) fu
 			log.Warnf("runner: %s cancelled in progress jobs during shutdown", resp.Msg.Runner.Name)
 		}
 
+		if poller.RegistrationRejected() {
+			removeEphemeralRegistration(reg, cfg.Runner.File)
+		}
+
 		if poller.Unregistered() {
 			return errors.New("runner is no longer registered with the server; please register it again")
 		}
@@ -252,6 +259,17 @@ func runDaemon(ctx context.Context, daemArgs *daemonArgs, configFile *string) fu
 	}
 }
 
+func removeEphemeralRegistration(reg *config.Registration, file string) {
+	if !reg.Ephemeral {
+		return
+	}
+	if err := os.Remove(file); err == nil {
+		log.Infof("removed the spent ephemeral registration %q", file)
+	} else if !os.IsNotExist(err) {
+		log.WithError(err).Warn("failed to remove the spent ephemeral registration")
+	}
+}
+
 type daemonArgs struct {
 	Once   bool
 	Labels string

internal/app/cmd/daemon_test.go

@@ -4,13 +4,24 @@
 package cmd
 
 import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
 	"testing"
+	"time"
 
 	"gitea.com/gitea/runner/internal/pkg/config"
 	"gitea.com/gitea/runner/internal/pkg/report"
 
+	"connectrpc.com/connect"
+	runnerv1 "gitea.dev/actionslib/runner/v1"
+	"gitea.dev/actionslib/runner/v1/runnerv1connect"
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 func TestResolveLabels(t *testing.T) {
@@ -70,3 +81,79 @@ func TestInitLoggingSetsLevelAndCaller(t *testing.T) {
 	// act plans a job on this logger, so a live task's secrets have to be masked out of it
 	require.IsType(t, report.MaskingFormatter(nil), log.StandardLogger().Formatter)
 }
+
+var errUnregisteredRunner = errors.New("rpc error: code = Unauthenticated desc = unregistered runner")
+
+type daemonTestService struct {
+	runnerv1connect.UnimplementedRunnerServiceHandler
+	declareErr, fetchErr error
+	beforeTask           bool
+	cancel               context.CancelFunc
+}
+
+func (s *daemonTestService) Declare(context.Context, *connect.Request[runnerv1.DeclareRequest]) (*connect.Response[runnerv1.DeclareResponse], error) {
+	if s.declareErr != nil {
+		return nil, s.declareErr
+	}
+	return connect.NewResponse(&runnerv1.DeclareResponse{Runner: &runnerv1.Runner{Name: "test"}}), nil
+}
+
+func (s *daemonTestService) FetchTask(context.Context, *connect.Request[runnerv1.FetchTaskRequest]) (*connect.Response[runnerv1.FetchTaskResponse], error) {
+	if s.fetchErr != nil {
+		return nil, s.fetchErr
+	}
+	if s.beforeTask {
+		s.cancel()
+		return nil, context.Canceled
+	}
+	return connect.NewResponse(&runnerv1.FetchTaskResponse{Task: &runnerv1.Task{Id: 1, Context: &structpb.Struct{}, WorkflowPayload: []byte("invalid: [")}}), nil
+}
+
+func (*daemonTestService) UpdateTask(_ context.Context, req *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error) {
+	return connect.NewResponse(&runnerv1.UpdateTaskResponse{State: req.Msg.State}), nil
+}
+
+func (*daemonTestService) UpdateLog(_ context.Context, req *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error) {
+	return connect.NewResponse(&runnerv1.UpdateLogResponse{AckIndex: req.Msg.Index + int64(len(req.Msg.Rows))}), nil
+}
+
+func TestDaemonRemovesOnlyConsumedEphemeralRegistration(t *testing.T) {
+	for _, tc := range []struct {
+		name                  string
+		ephemeral, beforeTask bool
+		declareErr, fetchErr  error
+		removed               bool
+	}{
+		{"ephemeral failed task", true, false, nil, nil, true},
+		{"ephemeral interrupted before task", true, true, nil, nil, false},
+		{"persistent once", false, false, nil, nil, false},
+		{"ephemeral rejected on fetch", true, false, nil, errUnregisteredRunner, true},
+		{"ephemeral proxy 401 on fetch", true, false, nil, connect.NewError(connect.CodeUnauthenticated, errors.New("HTTP status 401 Unauthorized")), false},
+		{"ephemeral rejected on declare", true, false, errUnregisteredRunner, nil, true},
+		{"persistent rejected on declare", false, false, errUnregisteredRunner, nil, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
+			defer cancel()
+			_, handler := runnerv1connect.NewRunnerServiceHandler(&daemonTestService{declareErr: tc.declareErr, fetchErr: tc.fetchErr, beforeTask: tc.beforeTask, cancel: cancel})
+			server := httptest.NewServer(http.StripPrefix("/api/actions", handler))
+			defer server.Close()
+			dir := t.TempDir()
+			regFile := filepath.Join(dir, "registration.json")
+			configFile := filepath.Join(dir, "config.yaml")
+			require.NoError(t, os.WriteFile(configFile, []byte("runner:\n  file: "+regFile+"\n  idle_cleanup_interval: 0s\ncache:\n  enabled: false\n"), 0o600))
+			require.NoError(t, config.SaveRegistration(regFile, &config.Registration{Address: server.URL, UUID: "test", Name: "test", Token: "test", Labels: []string{"host:host"}, Ephemeral: tc.ephemeral}))
+			err := runDaemon(ctx, &daemonArgs{Once: true}, &configFile)(nil, nil)
+			if tc.declareErr == nil && tc.fetchErr == nil {
+				require.NoError(t, err)
+			} else {
+				require.Error(t, err)
+			}
+			if tc.removed {
+				require.NoFileExists(t, regFile)
+			} else {
+				require.FileExists(t, regFile)
+			}
+		})
+	}
+}

internal/app/poll/poller.go

@@ -52,6 +52,8 @@ type Poller struct {
 
 	done chan struct{}
 
+	registrationRejected atomic.Bool
+
 	// unregistered is set when the server rejects the runner with an
 	// Unauthenticated response, meaning the runner is no longer registered.
 	unregistered       atomic.Bool
@@ -151,7 +153,7 @@ func (p *Poller) Poll() {
 	}
 }
 
-func (p *Poller) PollOnce() {
+func (p *Poller) PollOnce(onTaskFetched func()) {
 	defer close(p.done)
 	s := &workerState{}
 	for {
@@ -173,6 +175,7 @@ func (p *Poller) PollOnce() {
 			continue
 		}
 		s.resetBackoff()
+		onTaskFetched()
 		p.runTaskWithRecover(p.jobsCtx, task)
 		return
 	}
@@ -191,6 +194,11 @@ func (p *Poller) Unregistered() bool {
 	return p.unregistered.Load()
 }
 
+// RegistrationRejected reports whether polling stopped on an error matching client.IsRegistrationRejected.
+func (p *Poller) RegistrationRejected() bool {
+	return p.registrationRejected.Load()
+}
+
 // Ready reports whether the daemon can currently communicate with Gitea and
 // accept work, reusing the availability the poll loop last observed rather than
 // re-running the check. Transient transport failures are tolerated for grace.
@@ -368,9 +376,11 @@ func (p *Poller) fetchTask(ctx context.Context, s *workerState) (*runnerv1.Task,
 		// An Unauthenticated response means the server no longer knows this
 		// runner (e.g. it was deleted). Retrying forever is pointless, so stop
 		// polling and let the daemon exit with an error instead of spinning.
-		if isUnregistered(err) {
+		rejected := client.IsRegistrationRejected(err)
+		if rejected || connect.CodeOf(err) == connect.CodeUnauthenticated {
 			log.WithError(err).Error("server rejected the runner as unregistered, stopping poller")
 			p.unregistered.Store(true)
+			p.registrationRejected.Store(rejected)
 			p.shutdownPolling()
 			return nil, false
 		}
@@ -415,12 +425,6 @@ func (p *Poller) fetchTask(ctx context.Context, s *workerState) (*runnerv1.Task,
 	return resp.Msg.Task, true
 }
 
-func isUnregistered(err error) bool {
-	var connectErr *connect.Error
-	return errors.As(err, &connectErr) && (connectErr.Code() == connect.CodeUnauthenticated ||
-		connectErr.Code() == connect.CodeUnknown && connect.IsWireError(err) && connectErr.Message() == "rpc error: code = Unauthenticated desc = unregistered runner")
-}
-
 func (p *Poller) markHealthyPoll() {
 	p.lastHealthyPoll.Store(time.Now().UnixNano())
 	p.lastPollFailed.Store(false)

internal/app/poll/poller_test.go

@@ -106,13 +106,16 @@ func TestPoller_FetchTimeoutIsNoSignal(t *testing.T) {
 // response marks the runner as unregistered and cancels the polling context so
 // the daemon can exit instead of retrying forever.
 func TestPoller_FetchUnauthenticatedStopsPolling(t *testing.T) {
-	for name, fetchErr := range map[string]error{
-		"connect": connect_go.NewError(connect_go.CodeUnauthenticated, errors.New("unregistered runner")),
-		"gitea":   connect_go.NewWireError(connect_go.CodeUnknown, errors.New("rpc error: code = Unauthenticated desc = unregistered runner")),
+	for name, tc := range map[string]struct {
+		fetchErr error
+		rejected bool
+	}{
+		"connect": {connect_go.NewError(connect_go.CodeUnauthenticated, errors.New("unregistered runner")), false},
+		"gitea":   {connect_go.NewWireError(connect_go.CodeUnknown, errors.New("rpc error: code = Unauthenticated desc = unregistered runner")), true},
 	} {
 		t.Run(name, func(t *testing.T) {
 			client := mocks.NewClient(t)
-			client.On("FetchTask", mock.Anything, mock.Anything).Return(nil, fetchErr)
+			client.On("FetchTask", mock.Anything, mock.Anything).Return(nil, tc.fetchErr)
 
 			cfg, err := config.LoadDefault("")
 			require.NoError(t, err)
@@ -122,6 +125,7 @@ func TestPoller_FetchUnauthenticatedStopsPolling(t *testing.T) {
 			require.False(t, ok)
 
 			assert.True(t, p.Unregistered(), "runner should be marked unregistered")
+			assert.Equal(t, tc.rejected, p.RegistrationRejected())
 			assert.Equal(t, int64(0), s.consecutiveErrors, "unauthenticated must not drive error backoff")
 			select {
 			case <-p.pollingCtx.Done():
@@ -330,7 +334,7 @@ func TestPollerPollOnceCallsOnIdle(t *testing.T) {
 	poller := New(cfg, cli, runner)
 
 	var wg sync.WaitGroup
-	wg.Go(poller.PollOnce)
+	wg.Go(func() { poller.PollOnce(func() {}) })
 
 	require.Eventually(t, func() bool {
 		return runner.idleCalls.Load() > 0
@@ -342,6 +346,18 @@ func TestPollerPollOnceCallsOnIdle(t *testing.T) {
 	wg.Wait()
 }
 
+func TestPollerPollOnceCallsOnTaskFetchedBeforeRunningIt(t *testing.T) {
+	cli := mocks.NewClient(t)
+	cli.On("FetchTask", mock.Anything, mock.Anything).Return(connect_go.NewResponse(&runnerv1.FetchTaskResponse{Task: &runnerv1.Task{Id: 1}}), nil)
+	cfg, err := config.LoadDefault("")
+	require.NoError(t, err)
+	runner := &mockRunner{}
+	runsSeenByCallback := int64(-1)
+	New(cfg, cli, runner).PollOnce(func() { runsSeenByCallback = runner.running.Load() + runner.totalCompleted.Load() })
+	assert.Zero(t, runsSeenByCallback)
+	assert.Equal(t, int64(1), runner.totalCompleted.Load())
+}
+
 // TestPoller_ConcurrencyLimitedByCapacity verifies that with capacity=3 and
 // 6 available tasks, at most 3 tasks run concurrently, and FetchTask is
 // never called concurrently (single poller).

internal/pkg/client/client.go

@@ -5,6 +5,7 @@ package client
 
 import (
 	"context"
+	"errors"
 
 	"connectrpc.com/connect"
 	"gitea.dev/actionslib/runner/v1"
@@ -18,3 +19,10 @@ type Client interface {
 	UpdateLog(context.Context, *connect.Request[runnerv1.UpdateLogRequest]) (*connect.Response[runnerv1.UpdateLogResponse], error)
 	UpdateTask(context.Context, *connect.Request[runnerv1.UpdateTaskRequest]) (*connect.Response[runnerv1.UpdateTaskResponse], error)
 }
+
+// IsRegistrationRejected reports whether Gitea itself answered "unregistered runner", unlike a proxy's bare 401.
+func IsRegistrationRejected(err error) bool {
+	var connectErr *connect.Error
+	return errors.As(err, &connectErr) && connectErr.Code() == connect.CodeUnknown && connect.IsWireError(err) &&
+		connectErr.Message() == "rpc error: code = Unauthenticated desc = unregistered runner"
+}