uncloud/internal/machine/docker/service.go

package docker

import (
	"bufio"
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"strconv"
	"strings"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/containerd/errdefs"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/events"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/image"
	"github.com/docker/docker/client"
	"github.com/docker/docker/pkg/stdcopy"
	"github.com/jmoiron/sqlx"
	"github.com/psviderski/uncloud/pkg/api"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// Service provides higher-level Docker operations that extends Docker API with Uncloud-specific data
// from the machine database.
type Service struct {
	// Client is a Docker client for managing Docker resources.
	Client client.APIClient
	// db is a connection to the machine database.
	db *sqlx.DB
	// containersChanged signals that service containers changed in the machine database without a corresponding
	// Docker event, for example, when a container spec is updated in place.
	containersChanged chan struct{}
}

// NewService creates a new Docker service instance.
func NewService(client client.APIClient, db *sqlx.DB) *Service {
	return &Service{
		Client: client,
		db:     db,
		// Buffering as a pending signal is enough for the consumer to pick up all changes made before it reacts.
		containersChanged: make(chan struct{}, 1),
	}
}

// NotifyContainersChanged signals that service containers changed in the machine database. It doesn't block
// if a signal is already pending.
func (s *Service) NotifyContainersChanged() {
	select {
	case s.containersChanged <- struct{}{}:
	default:
	}
}

// ContainersChanged returns a channel that receives a signal when service containers change in the machine database
// without a corresponding Docker event.
func (s *Service) ContainersChanged() <-chan struct{} {
	return s.containersChanged
}

// WatchDaemonRestart watches the Docker engine and signals on the returned channel whenever the daemon
// becomes reachable again after a disconnect, which typically indicates a restart.
func (s *Service) WatchDaemonRestart(ctx context.Context) <-chan struct{} {
	// Buffering as a pending signal is enough for the consumer to react to a restart.
	ch := make(chan struct{}, 1)

	go func() {
		boff := backoff.WithContext(backoff.NewExponentialBackOff(
			backoff.WithInitialInterval(1*time.Second),
			backoff.WithMaxInterval(30*time.Second),
			backoff.WithMaxElapsedTime(0),
		), ctx)

		// firstConnect tracks the initial connection so it is not reported as a restart.
		firstConnect := true

		watch := func() error {
			// Confirm the daemon is reachable. While it is down this fails and is retried with backoff.
			if _, err := s.Client.Ping(ctx); err != nil {
				firstConnect = false
				return fmt.Errorf("ping Docker daemon: %w", err)
			}
			// Reset the backoff after a successful connection so the next reconnect retries quickly.
			boff.Reset()

			if firstConnect {
				firstConnect = false
			} else {
				// The daemon became reachable again after a disconnect. It likely restarted.
				slog.Debug("Docker daemon reconnected.")
				select {
				case ch <- struct{}{}:
				default:
				}
			}

			// Subscribe to daemon events and block until the stream breaks (daemon stopped/restarting) or
			// the context is done. Filter to non-existent events to not receive any events.
			eventCh, errCh := s.Client.Events(ctx, events.ListOptions{
				Filters: filters.NewArgs(filters.Arg("type", "invalid")),
			})
			for {
				select {
				case <-ctx.Done():
					return nil
				case <-eventCh:
					// This shouldn't emit anything but drain just in case. We only care about the stream lifecycle.
				case streamErr := <-errCh:
					if ctx.Err() != nil {
						return nil
					}
					return fmt.Errorf("docker event stream closed: %w", streamErr)
				}
			}
		}

		// Retry watching the engine until the context is cancelled.
		_ = backoff.Retry(watch, boff)
	}()

	return ch
}

// InspectServiceContainer inspects a Docker container and retrieves its associated ServiceSpec
// from the machine database, returning a complete ServiceContainer.
func (s *Service) InspectServiceContainer(ctx context.Context, nameOrID string) (api.ServiceContainer, error) {
	var serviceCtr api.ServiceContainer

	ctr, err := s.Client.ContainerInspect(ctx, nameOrID)
	if err != nil {
		return serviceCtr, err
	}
	if _, ok := ctr.Config.Labels[api.LabelManaged]; !ok {
		return serviceCtr, fmt.Errorf("container '%s' is not managed by Uncloud", nameOrID)
	}

	serviceCtr.Container = api.Container{InspectResponse: ctr}

	// Retrieve ServiceSpec from the machine database.
	var specBytes []byte
	err = s.db.QueryRowContext(ctx, `SELECT service_spec FROM containers WHERE id = $1`, ctr.ID).Scan(&specBytes)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			// If this happens, there is a bug in the code, or someone manually removed the container from the DB,
			// or created a managed container out of band or by previous uncloud installation.
			return serviceCtr, fmt.Errorf("service spec not found for container '%s' in machine DB", ctr.ID)
		}
		return serviceCtr, fmt.Errorf("get service spec for container '%s' from machine DB: %w", ctr.ID, err)
	}

	if err = json.Unmarshal(specBytes, &serviceCtr.ServiceSpec); err != nil {
		return serviceCtr, fmt.Errorf("unmarshal service spec for container '%s': %w", ctr.ID, err)
	}

	return serviceCtr, nil
}

// UpdateServiceContainer updates the requested properties of a service container in place without recreating it.
// It applies the CPU and memory limits to the container and updates the service spec stored for the container.
func (s *Service) UpdateServiceContainer(ctx context.Context, id string, opts api.UpdateContainerOptions) error {
	if opts.ServiceLabels != nil {
		if err := api.ValidateLabelKeys(*opts.ServiceLabels); err != nil {
			return errdefs.ErrInvalidArgument.WithMessage(fmt.Sprintf("invalid service labels: %v", err))
		}
	}

	ctr, err := s.InspectServiceContainer(ctx, id)
	if err != nil {
		return err
	}
	spec := ctr.ServiceSpec

	resources := spec.Container.Resources
	if opts.CPU != nil {
		resources.CPU = *opts.CPU
	}
	if opts.Memory != nil {
		resources.Memory = *opts.Memory
	}
	if opts.MemoryReservation != nil {
		resources.MemoryReservation = *opts.MemoryReservation
	}
	if resources.CPU < 0 || resources.Memory < 0 || resources.MemoryReservation < 0 {
		return errdefs.ErrInvalidArgument.WithMessage("CPU and memory limits must not be negative")
	}
	// Compare with the limits Docker applies to the container rather than the stored spec. They may differ if
	// a previous update applied the limits but failed to store the spec, or the container was updated outside
	// of Uncloud.
	applied := ctr.AppliedResources()
	if err = applied.ValidateUpdateInPlace(resources); err != nil {
		return errdefs.ErrFailedPrecondition.WithMessage(
			fmt.Sprintf("container '%s' can't be updated without recreating it: %v", ctr.ID, err))
	}

	// Only CPU and memory limits can differ after the in-place update check above.
	if !applied.Equals(resources) {
		updateResources := container.Resources{
			NanoCPUs:          resources.CPU,
			Memory:            resources.Memory,
			MemoryReservation: resources.MemoryReservation,
		}
		// Docker defaults the swap limit to twice the memory limit when creating a container and rejects raising
		// the memory limit above the current swap limit, so update both to keep the same semantics.
		if resources.Memory > 0 {
			updateResources.MemorySwap = 2 * resources.Memory
		}
		if _, err = s.Client.ContainerUpdate(ctx, ctr.ID, container.UpdateConfig{
			Resources: updateResources,
		}); err != nil {
			return fmt.Errorf("update container resources: %w", err)
		}
	}
	spec.Container.Resources = resources

	if opts.ServiceLabels != nil {
		spec.Labels = nil
		if len(*opts.ServiceLabels) > 0 {
			spec.Labels = *opts.ServiceLabels
		}
	}

	specBytes, err := json.Marshal(spec)
	if err != nil {
		return fmt.Errorf("marshal service spec: %w", err)
	}
	res, err := s.db.ExecContext(ctx,
		`UPDATE containers SET service_spec = $1, updated_at = datetime('subsecond') WHERE id = $2`,
		string(specBytes), ctr.ID)
	if err != nil {
		return fmt.Errorf("update container in database: %w", err)
	}
	if n, err := res.RowsAffected(); err == nil && n == 0 {
		return errdefs.ErrNotFound.WithMessage(fmt.Sprintf("container '%s' not found in machine database", ctr.ID))
	}

	// Propagate the updated spec to the cluster store without waiting for the next regular sync.
	s.NotifyContainersChanged()

	return nil
}

// ListServiceContainersResult holds the result of listing service containers, split into regular
// service containers and one-off hook containers.
type ListServiceContainersResult struct {
	Containers     []api.ServiceContainer
	HookContainers []api.ServiceContainer
}

// ListServiceContainers lists Docker containers that belong to the service with the given name or ID.
// If serviceIDOrName is empty, all service containers are returned. The opts parameter allows additional filtering.
func (s *Service) ListServiceContainers(
	ctx context.Context, serviceNameOrID string, opts container.ListOptions,
) (ListServiceContainersResult, error) {
	var result ListServiceContainersResult

	if opts.Filters.Len() == 0 {
		opts.Filters = filters.NewArgs()
	}
	// Add labels to existing filters to list only Uncloud-managed service containers.
	opts.Filters.Add("label", api.LabelServiceID)
	opts.Filters.Add("label", api.LabelManaged)

	containerSummaries, err := s.Client.ContainerList(ctx, opts)
	if err != nil {
		return result, err
	}

	for _, cs := range containerSummaries {
		// Filter by service name or ID if provided.
		if serviceNameOrID != "" &&
			cs.Labels[api.LabelServiceID] != serviceNameOrID &&
			cs.Labels[api.LabelServiceName] != serviceNameOrID {
			continue
		}

		ctr, err := s.InspectServiceContainer(ctx, cs.ID)
		if err != nil {
			// Log error but continue with other containers.
			slog.Error("Failed to inspect service container.", "service", serviceNameOrID, "id", cs.ID, "err", err)
			continue
		}

		if ctr.IsHook() {
			result.HookContainers = append(result.HookContainers, ctr)
		} else {
			result.Containers = append(result.Containers, ctr)
		}
	}

	return result, nil
}

// IsContainerdImageStoreEnabled checks if Docker is configured to use the containerd image store:
// https://docs.docker.com/engine/storage/containerd/
func (s *Service) IsContainerdImageStoreEnabled(ctx context.Context) (bool, error) {
	info, err := s.Client.Info(ctx)
	if err != nil {
		return false, fmt.Errorf("get Docker info: %w", err)
	}

	return strings.Contains(fmt.Sprintf("%s", info.DriverStatus), "containerd.snapshotter"), nil
}

type Images struct {
	// Images is a list of images present in the Docker image store (either internal or containerd).
	Images []image.Summary
	// ContainerdStore indicates whether Docker is using the containerd image store.
	ContainerdStore bool
}

// ListImages lists Docker images with the given options and indicates whether Docker is using the containerd
// image store. It always includes image manifests in the response if the store is containerd.
func (s *Service) ListImages(ctx context.Context, opts image.ListOptions) (Images, error) {
	var imagesResp Images

	// Always include the image manifests in the response.
	opts.Manifests = true
	images, err := s.Client.ImageList(ctx, opts)
	if err != nil {
		return imagesResp, status.Errorf(codes.Internal, "list images: %v", err)
	}

	isContainerdStore, err := s.IsContainerdImageStoreEnabled(ctx)
	if err != nil {
		return imagesResp, status.Errorf(codes.Internal, "check if Docker uses containerd image store: %v", err)
	}

	imagesResp = Images{
		Images:          images,
		ContainerdStore: isContainerdStore,
	}

	return imagesResp, nil
}

// ContainerLogs streams logs from a container and returns entries via a channel.
// The channel is closed when streaming completes or context is cancelled.
func (s *Service) ContainerLogs(
	ctx context.Context, containerID string, opts api.ServiceLogsOptions,
) (<-chan api.LogEntry, error) {
	ctr, err := s.Client.ContainerInspect(ctx, containerID)
	if err != nil {
		return nil, fmt.Errorf("inspect container '%s': %w", containerID, err)
	}
	isTTY := ctr.Config != nil && ctr.Config.Tty

	dockerOpts := container.LogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Follow:     opts.Follow,
		Tail:       strconv.FormatInt(int64(opts.Tail), 10),
		Since:      opts.Since,
		Until:      opts.Until,
		Timestamps: true,
	}

	reader, err := s.Client.ContainerLogs(ctx, containerID, dockerOpts)
	if err != nil {
		return nil, err
	}

	outCh := make(chan api.LogEntry)
	stdoutWriter := &logsChannelWriter{ctx: ctx, ch: outCh, isStderr: false}
	stderrWriter := &logsChannelWriter{ctx: ctx, ch: outCh, isStderr: true}

	// Wrap the context in a cancellable one to unblock the second goroutine when log copying completes.
	ctx, cancel := context.WithCancel(ctx)

	// Copy logs in a goroutine to be able to handle context cancellation.
	go func() {
		defer close(outCh)
		defer cancel()

		// Docker returns raw stdout for TTY containers and multiplexed stdout/stderr otherwise.
		var err error
		if isTTY {
			_, err = copyRawContainerLogs(stdoutWriter, reader)
		} else {
			_, err = stdcopy.StdCopy(stdoutWriter, stderrWriter, reader)
		}
		if err != nil {
			// Send error as the last entry.
			select {
			case outCh <- api.LogEntry{Err: fmt.Errorf("copy container logs: %w", err)}:
			case <-ctx.Done():
			}
		}
	}()

	// Close the reader when the context is done to cancel log copying if it's still running.
	go func() {
		<-ctx.Done()
		reader.Close()
	}()

	return outCh, nil
}

// copyRawContainerLogs copies a raw TTY log stream one line at a time so each write produces one log entry.
func copyRawContainerLogs(dst io.Writer, src io.Reader) (written int64, _ error) {
	reader := bufio.NewReader(src)
	for {
		line, readErr := reader.ReadBytes('\n')
		if len(line) > 0 {
			n, writeErr := dst.Write(line)
			written += int64(n)
			if writeErr != nil {
				return written, writeErr
			}
			if n != len(line) {
				return written, io.ErrShortWrite
			}
		}

		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return written, nil
			}
			return written, readErr
		}
	}
}

// logsChannelWriter sends container log writes to a channel.
type logsChannelWriter struct {
	ctx      context.Context
	ch       chan<- api.LogEntry
	isStderr bool
}

func (w *logsChannelWriter) Write(data []byte) (n int, err error) {
	// Parse timestamp and message from the demultiplexed Docker log payload if the data looks like it contains one.
	// Format: 2025-01-01T00:00:00.000000000Z message
	timestamp := time.Time{}
	message := data
	if len(data) > 30 && data[4] == '-' && data[7] == '-' && data[10] == 'T' {
		timestampPart, messagePart, found := bytes.Cut(data, []byte(" "))
		if found {
			timestamp, err = time.Parse(time.RFC3339Nano, string(timestampPart))
			if err != nil {
				timestamp = time.Time{}
			}
			message = messagePart
		}
	}

	entry := api.LogEntry{
		Timestamp: timestamp,
		// Clone is required because message is a slice into data, which stdcopy.StdCopy may reuse
		// after Write returns but before the entry is consumed from the channel.
		Message: bytes.Clone(message),
	}
	if w.isStderr {
		entry.Stream = api.LogStreamStderr
	} else {
		entry.Stream = api.LogStreamStdout
	}

	select {
	case w.ch <- entry:
		return len(data), nil
	case <-w.ctx.Done():
		return 0, w.ctx.Err()
	}
}