package-updater/internal/agent/updater/selfupdate.go

package updater

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"strings"
	"time"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/pkg/stdcopy"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
)

// Labels on a self-update helper container. The new agent finds the helper
// by them, and the inventory skips containers that carry the job label.
const (
	labelSelfUpdateJob     = "se.shcizo.package-updater.self-update.job"
	labelSelfUpdateRef     = "se.shcizo.package-updater.self-update.ref"
	labelSelfUpdateStarted = "se.shcizo.package-updater.self-update.started"
	labelSelfUpdateProject = "se.shcizo.package-updater.self-update.project"
	labelSelfUpdateService = "se.shcizo.package-updater.self-update.service"
	labelSelfUpdateOutput  = "se.shcizo.package-updater.self-update.output"
)

// selfUpdateOutputMax caps the job output stored in the helper container's
// output label. The old agent's output so far is all the new agent has of it.
const selfUpdateOutputMax = 32 << 10

// truncatedMarker starts the stored output when its head was cut off.
const truncatedMarker = "…truncated\n"

// resumeSeparator follows the old agent's output in a resumed job's output.
const resumeSeparator = "--- agent restarted; output of the update helper ---\n"

// HelperDocker is the subset of the Docker SDK the self-update helper
// container needs. *client.Client satisfies it.
type HelperDocker interface {
	ContainerInspect(ctx context.Context, containerID string) (container.InspectResponse, error)
	ContainerCreate(ctx context.Context, config *container.Config, hostConfig *container.HostConfig, networkingConfig *network.NetworkingConfig, platform *ocispec.Platform, containerName string) (container.CreateResponse, error)
	ContainerStart(ctx context.Context, containerID string, options container.StartOptions) error
	ContainerWait(ctx context.Context, containerID string, condition container.WaitCondition) (<-chan container.WaitResponse, <-chan error)
	ContainerLogs(ctx context.Context, containerID string, options container.LogsOptions) (io.ReadCloser, error)
	ContainerRemove(ctx context.Context, containerID string, options container.RemoveOptions) error
}

// selfUpdate runs `up -d` for the agent's own service in a short-lived helper
// container. Recreating the service stops this process, so the helper, made
// from the agent's own image with the agent's own mounts, does it and
// outlives the agent. args are the complete `docker` arguments of the call.
//
// If the helper exits while this agent still runs, selfUpdate copies its logs
// into the job output, removes it and returns the outcome. If ctx ends first
// (the agent is being replaced, or the job timed out) it returns ctx's error
// and leaves the helper to the new agent, see ResumeSelfUpdates.
func (e *ComposeExecutor) selfUpdate(ctx context.Context, job Job, self container.InspectResponse, workingDir string, args []string, pinned string, out io.Writer, soFar *tailBuffer) (string, error) {
	fmt.Fprintf(out, "$ docker %s (in an update helper container)\n", strings.Join(args, " "))
	cfg := &container.Config{
		Image:      self.Image,
		Entrypoint: []string{"docker"},
		Cmd:        args,
		WorkingDir: workingDir,
		Env:        dockerEnv(),
		Labels: map[string]string{
			labelSelfUpdateJob:     job.ID,
			labelSelfUpdateRef:     pinned,
			labelSelfUpdateStarted: job.Started.UTC().Format(time.RFC3339Nano),
			labelSelfUpdateProject: job.ComposeProject,
			labelSelfUpdateService: job.ComposeService,
			labelSelfUpdateOutput:  outputLabel(soFar),
		},
	}
	created, err := e.cli.ContainerCreate(ctx, cfg, &container.HostConfig{Mounts: helperMounts(self)}, nil, nil, "")
	if err != nil {
		return "", fmt.Errorf("create update helper: %w", err)
	}
	// Wait is registered before the start so a fast exit is not missed.
	waitCh, errCh := e.cli.ContainerWait(ctx, created.ID, container.WaitConditionNextExit)
	if err := e.cli.ContainerStart(ctx, created.ID, container.StartOptions{}); err != nil {
		if rmErr := e.cli.ContainerRemove(context.WithoutCancel(ctx), created.ID, container.RemoveOptions{Force: true}); rmErr != nil {
			fmt.Fprintf(out, "warning: remove the update helper: %v\n", rmErr)
		}
		return "", fmt.Errorf("start update helper: %w", err)
	}
	code, err := awaitExit(ctx, waitCh, errCh)
	if err != nil {
		return "", err
	}

	tail := &tailBuffer{max: errorOutputTail}
	if err := finishHelper(ctx, e.cli, created.ID, io.MultiWriter(out, tailWriter{tail})); err != nil {
		fmt.Fprintf(out, "warning: %v\n", err)
	}
	if code != 0 {
		return "", exitError(code, tail)
	}
	return pinned, nil
}

// outputLabel returns the job output so far as a label value.
func outputLabel(b *tailBuffer) string {
	s := strings.ToValidUTF8(b.String(), "")
	if b.Truncated() {
		s = truncatedMarker + s
	}
	return s
}

// helperMounts copies the agent's own bind and volume mounts: the Docker
// socket, the stack directories and the registry credentials the compose call
// needs are all among them.
func helperMounts(self container.InspectResponse) []mount.Mount {
	var mounts []mount.Mount
	for _, m := range self.Mounts {
		switch m.Type {
		case mount.TypeBind:
			mounts = append(mounts, mount.Mount{Type: mount.TypeBind, Source: m.Source, Target: m.Destination, ReadOnly: !m.RW})
		case mount.TypeVolume:
			mounts = append(mounts, mount.Mount{Type: mount.TypeVolume, Source: m.Name, Target: m.Destination, ReadOnly: !m.RW})
		}
	}
	return mounts
}

// awaitExit waits for a ContainerWait result and returns the exit code. When
// ctx ends it returns ctx's error.
func awaitExit(ctx context.Context, waitCh <-chan container.WaitResponse, errCh <-chan error) (int64, error) {
	select {
	case r := <-waitCh:
		if r.Error != nil {
			return 0, fmt.Errorf("wait for update helper: %s", r.Error.Message)
		}
		return r.StatusCode, nil
	case err := <-errCh:
		if ctx.Err() != nil {
			return 0, ctx.Err()
		}
		return 0, fmt.Errorf("wait for update helper: %w", err)
	case <-ctx.Done():
		return 0, ctx.Err()
	}
}

// finishHelper copies the exited helper's logs to w and removes the helper.
// Both are attempted; the first failure is returned.
func finishHelper(ctx context.Context, cli HelperDocker, id string, w io.Writer) error {
	var first error
	rc, err := cli.ContainerLogs(ctx, id, container.LogsOptions{ShowStdout: true, ShowStderr: true})
	if err == nil {
		_, err = stdcopy.StdCopy(w, w, rc)
		rc.Close()
	}
	if err != nil {
		first = fmt.Errorf("read the update helper's logs: %w", err)
	}
	if err := cli.ContainerRemove(ctx, id, container.RemoveOptions{Force: true}); err != nil && first == nil {
		first = fmt.Errorf("remove the update helper: %w", err)
	}
	return first
}

// exitError describes a failed helper like a failed `docker compose` call.
func exitError(code int64, tail *tailBuffer) error {
	return composeError("up", fmt.Errorf("exit status %d", code), tail)
}

// ResumeDocker is the subset of the Docker SDK ResumeSelfUpdates needs.
// *client.Client satisfies it.
type ResumeDocker interface {
	SelfInspector
	HelperDocker
	ContainerList(ctx context.Context, opts container.ListOptions) ([]container.Summary, error)
}

// ResumeSelfUpdates finishes the jobs a previous agent process started with
// selfUpdate. It finds the helper containers made for this agent's own
// Compose service and registers each with runner.Resume, which reports the
// job under its old ID once the helper has exited. Call it before the runner
// and the connection start, so the first hello already lists these jobs. A
// helper with unusable labels is resumed too; its job fails with an error
// that says so.
func ResumeSelfUpdates(ctx context.Context, cli ResumeDocker, selfID string, runner *Runner, log *slog.Logger) error {
	self, err := inspectSelf(ctx, cli, selfID)
	if err != nil {
		return err
	}
	if self.Config == nil {
		return nil
	}
	project, service := self.Config.Labels[labelComposeProject], self.Config.Labels[labelComposeService]
	if project == "" || service == "" {
		return nil
	}
	cts, err := cli.ContainerList(ctx, container.ListOptions{
		All:     true,
		Filters: filters.NewArgs(filters.Arg("label", labelSelfUpdateJob)),
	})
	if err != nil {
		return fmt.Errorf("list update helpers: %w", err)
	}
	for _, ct := range cts {
		if ct.Labels[labelSelfUpdateProject] != project || ct.Labels[labelSelfUpdateService] != service {
			continue
		}
		jobID := ct.Labels[labelSelfUpdateJob]
		if jobID == "" {
			log.Warn("self_update_helper_without_job", "helper", ct.ID)
			continue
		}
		ref, started, labelErr := parseHelperLabels(ct.Labels)
		if started.IsZero() {
			started = time.Now()
		}
		log.Info("self_update_resumed", "job_id", jobID, "helper", ct.ID)
		runner.Resume(jobID, started, resumeFunc(cli, ct.ID, ct.Labels[labelSelfUpdateOutput], ref, labelErr))
	}
	return nil
}

// parseHelperLabels reads the pinned ref and start time from a helper's
// labels. started is zero when it is missing or unreadable, and err then says
// what is wrong.
func parseHelperLabels(labels map[string]string) (ref string, started time.Time, err error) {
	ref = labels[labelSelfUpdateRef]
	if ref == "" {
		err = errors.New("update helper has no ref label")
	}
	raw := labels[labelSelfUpdateStarted]
	if t, perr := time.Parse(time.RFC3339Nano, raw); perr == nil {
		started = t
	} else if err == nil {
		err = fmt.Errorf("update helper has no usable started label %q", raw)
	}
	return ref, started, err
}

// resumeFunc waits for the helper to stop, then reports the old agent's
// output, a separator line and the helper's own logs as the job output,
// removes the helper and returns the pinned ref, or an error when the helper
// failed or its labels were unusable.
func resumeFunc(cli HelperDocker, helperID, earlier, ref string, labelErr error) ResumeFunc {
	return func(ctx context.Context, out io.Writer) (string, error) {
		// A helper that was created but never started would read as exited
		// with code 0, so it must not be mistaken for a finished update.
		if insp, err := cli.ContainerInspect(ctx, helperID); err == nil && insp.ContainerJSONBase != nil &&
			insp.State != nil && insp.State.Status == container.StateCreated {
			if err := cli.ContainerRemove(context.WithoutCancel(ctx), helperID, container.RemoveOptions{Force: true}); err != nil {
				fmt.Fprintf(out, "warning: remove the update helper: %v\n", err)
			}
			return "", errors.New("update helper was never started")
		}
		waitCh, errCh := cli.ContainerWait(ctx, helperID, container.WaitConditionNotRunning)
		code, err := awaitExit(ctx, waitCh, errCh)
		if err != nil {
			return "", err
		}
		io.WriteString(out, earlier)
		if earlier != "" && !strings.HasSuffix(earlier, "\n") {
			io.WriteString(out, "\n")
		}
		io.WriteString(out, resumeSeparator)
		tail := &tailBuffer{max: errorOutputTail}
		if err := finishHelper(ctx, cli, helperID, io.MultiWriter(out, tailWriter{tail})); err != nil {
			fmt.Fprintf(out, "warning: %v\n", err)
		}
		switch {
		case labelErr != nil:
			return "", labelErr
		case code != 0:
			return "", exitError(code, tail)
		}
		return ref, nil
	}
}