From e88f105490a335e69698e0b7d05535356c0fb85c Mon Sep 17 00:00:00 2001
From: Samuel Enocsson <samuel.enocsson@gmail.com>
Date: Fri, 22 May 2026 12:06:07 +0200
Subject: [PATCH] feat(selfupdate): defer self-replacement until response
 flushed

---
 internal/selfupdate/selfupdate.go      | 56 +++++++++++++++++
 internal/selfupdate/selfupdate_test.go | 87 ++++++++++++++++++++++++++
 2 files changed, 143 insertions(+)
 create mode 100644 internal/selfupdate/selfupdate.go
 create mode 100644 internal/selfupdate/selfupdate_test.go

diff --git a/internal/selfupdate/selfupdate.go b/internal/selfupdate/selfupdate.go
new file mode 100644
index 0000000..e901d03
--- /dev/null
+++ b/internal/selfupdate/selfupdate.go
@@ -0,0 +1,56 @@
+// Package selfupdate handles the special case where the update target
+// is the running service's own container. We must finish writing the
+// HTTP response (and flush + close the connection) before exec'ing
+// `docker compose up -d` against ourselves, otherwise the response is
+// lost when the container is replaced.
+package selfupdate
+
+import (
+	"context"
+	"time"
+
+	"github.com/shcizo/package-updater/internal/discovery"
+)
+
+// IsSelf reports whether job targets the running service.
+// Matches on Compose project name (which is also typically the
+// service name for single-service stacks).
+func IsSelf(job discovery.Job, selfProject string) bool {
+	return job.Project == selfProject
+}
+
+// innerExec is the executor abstraction we wrap.
+type innerExec interface {
+	Execute(ctx context.Context, job discovery.Job) error
+}
+
+// Wrapped wraps an Executor with self-update-aware deferred execution.
+type Wrapped struct {
+	inner       innerExec
+	selfProject string
+	delay       time.Duration
+}
+
+// Wrap returns a Wrapped that defers exec until after flush() for
+// self-updates. The delay is added after flush before exec, so the
+// kernel TCP buffer has time to drain.
+func Wrap(inner innerExec, selfProject string, delay time.Duration) *Wrapped {
+	return &Wrapped{inner: inner, selfProject: selfProject, delay: delay}
+}
+
+// ExecuteWithFlush runs job, invoking flush() before exec for self-updates
+// and waiting `delay` after flush. For non-self jobs, exec happens
+// normally and flush is not invoked at all (the HTTP layer decides
+// when to flush in that case).
+func (w *Wrapped) ExecuteWithFlush(ctx context.Context, job discovery.Job, flush func()) error {
+	if !IsSelf(job, w.selfProject) {
+		return w.inner.Execute(ctx, job)
+	}
+	flush()
+	select {
+	case <-time.After(w.delay):
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+	return w.inner.Execute(ctx, job)
+}
diff --git a/internal/selfupdate/selfupdate_test.go b/internal/selfupdate/selfupdate_test.go
new file mode 100644
index 0000000..121181c
--- /dev/null
+++ b/internal/selfupdate/selfupdate_test.go
@@ -0,0 +1,87 @@
+package selfupdate_test
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/shcizo/package-updater/internal/discovery"
+	"github.com/shcizo/package-updater/internal/selfupdate"
+	"github.com/stretchr/testify/require"
+)
+
+type recExec struct {
+	mu     sync.Mutex
+	called bool
+	at     time.Time
+}
+
+func (r *recExec) Execute(_ context.Context, _ discovery.Job) error {
+	r.mu.Lock()
+	r.called = true
+	r.at = time.Now()
+	r.mu.Unlock()
+	return nil
+}
+
+func TestIsSelf(t *testing.T) {
+	require.True(t, selfupdate.IsSelf(
+		discovery.Job{Project: "package-updater", Service: "package-updater"},
+		"package-updater",
+	))
+	require.False(t, selfupdate.IsSelf(
+		discovery.Job{Project: "other", Service: "web"},
+		"package-updater",
+	))
+}
+
+func TestWrap_DefersSelf(t *testing.T) {
+	inner := &recExec{}
+	wrapped := selfupdate.Wrap(inner, "package-updater", 30*time.Millisecond)
+
+	flushed := make(chan time.Time, 1)
+	flush := func() { flushed <- time.Now() }
+
+	job := discovery.Job{Project: "package-updater", Service: "package-updater"}
+	require.NoError(t, wrapped.ExecuteWithFlush(context.Background(), job, flush))
+
+	flushAt := <-flushed
+	inner.mu.Lock()
+	require.True(t, inner.called)
+	require.True(t, inner.at.After(flushAt))
+	inner.mu.Unlock()
+}
+
+func TestWrap_NormalJobIsImmediate(t *testing.T) {
+	inner := &recExec{}
+	wrapped := selfupdate.Wrap(inner, "package-updater", 30*time.Millisecond)
+
+	flushed := make(chan time.Time, 1)
+	flush := func() { flushed <- time.Now() }
+
+	job := discovery.Job{Project: "other", Service: "web"}
+	require.NoError(t, wrapped.ExecuteWithFlush(context.Background(), job, flush))
+
+	inner.mu.Lock()
+	require.True(t, inner.called)
+	inner.mu.Unlock()
+	select {
+	case <-flushed:
+		t.Fatal("flush should not be called for non-self jobs")
+	default:
+	}
+}
+
+func TestWrap_PropagatesError(t *testing.T) {
+	wrapped := selfupdate.Wrap(failExec{}, "x", 1*time.Millisecond)
+	err := wrapped.ExecuteWithFlush(context.Background(), discovery.Job{Project: "other"}, func() {})
+	require.Error(t, err)
+}
+
+type failExec struct{}
+
+func (failExec) Execute(_ context.Context, _ discovery.Job) error {
+	return errors.New("boom")
+}