run/spool/artifact_queue.go

package spool

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"browser.local/run/protocol"
)

type ArtifactQueue struct {
	dir string
}

func NewArtifactQueue(dir string) (ArtifactQueue, error) {
	if strings.TrimSpace(dir) == "" {
		return ArtifactQueue{}, fmt.Errorf("spool directory is required")
	}
	artifactDir := filepath.Join(dir, "artifacts")
	if err := os.MkdirAll(artifactDir, 0o755); err != nil {
		return ArtifactQueue{}, fmt.Errorf("create artifact queue: %w", err)
	}
	return ArtifactQueue{dir: artifactDir}, nil
}

func (queue ArtifactQueue) Enqueue(chunk protocol.ArtifactChunkUploadRequest) error {
	path := queue.chunkPath(chunk)
	if existing, err := readArtifactChunk(path); err == nil {
		if existing.TransferID == chunk.TransferID && existing.ArtifactID == chunk.ArtifactID && existing.ChunkIndex == chunk.ChunkIndex && existing.Checksum == chunk.Checksum {
			return nil
		}
		return fmt.Errorf("artifact queue chunk conflicts with committed chunk")
	} else if !os.IsNotExist(err) {
		return err
	}
	payloadPath := artifactChunkPayloadPath(path)
	if err := writeArtifactQueueFile(payloadPath, chunk.Payload); err != nil {
		return err
	}
	metadata := chunk
	metadata.Payload = nil
	var body bytes.Buffer
	if err := json.NewEncoder(&body).Encode(metadata); err != nil {
		_ = os.Remove(payloadPath)
		return fmt.Errorf("encode artifact queue chunk metadata: %w", err)
	}
	if err := writeArtifactQueueFile(path, body.Bytes()); err != nil {
		_ = os.Remove(payloadPath)
		return err
	}
	return nil
}

type ArtifactChunkClient interface {
	UploadArtifactChunk(context.Context, protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error)
}

func (queue ArtifactQueue) Flush(ctx context.Context, client ArtifactChunkClient) (int, error) {
	if client == nil {
		return 0, fmt.Errorf("artifact chunk client is required")
	}
	pending, err := queue.Pending()
	if err != nil {
		return 0, err
	}
	acknowledged := 0
	for _, chunk := range pending {
		if err := ctx.Err(); err != nil {
			return acknowledged, err
		}
		response, err := client.UploadArtifactChunk(ctx, chunk)
		if err != nil {
			return acknowledged, err
		}
		if !response.Accepted || response.TransferID != chunk.TransferID || response.ArtifactID != chunk.ArtifactID || response.ChunkIndex != chunk.ChunkIndex {
			return acknowledged, fmt.Errorf("platform artifact acknowledgement does not match pending chunk")
		}
		if err := queue.Ack(response); err != nil {
			return acknowledged, err
		}
		acknowledged++
	}
	return acknowledged, nil
}

func readArtifactChunk(path string) (protocol.ArtifactChunkUploadRequest, error) {
	file, err := os.Open(path)
	if err != nil {
		return protocol.ArtifactChunkUploadRequest{}, err
	}
	defer file.Close()
	var chunk protocol.ArtifactChunkUploadRequest
	if err := json.NewDecoder(file).Decode(&chunk); err != nil {
		return protocol.ArtifactChunkUploadRequest{}, fmt.Errorf("decode artifact queue chunk: %w", err)
	}
	payload, err := os.ReadFile(artifactChunkPayloadPath(path))
	if err != nil {
		return protocol.ArtifactChunkUploadRequest{}, fmt.Errorf("read artifact queue chunk payload: %w", err)
	}
	chunk.Payload = payload
	return chunk, nil
}

func (queue ArtifactQueue) Pending() ([]protocol.ArtifactChunkUploadRequest, error) {
	entries, err := os.ReadDir(queue.dir)
	if err != nil {
		return nil, fmt.Errorf("read artifact queue: %w", err)
	}
	paths := make([]string, 0, len(entries))
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}
		paths = append(paths, filepath.Join(queue.dir, entry.Name()))
	}
	sort.Strings(paths)
	chunks := make([]protocol.ArtifactChunkUploadRequest, 0, len(paths))
	for _, path := range paths {
		chunk, err := readArtifactChunk(path)
		if err != nil {
			return nil, err
		}
		chunks = append(chunks, chunk)
	}
	return chunks, nil
}

func (queue ArtifactQueue) Ack(response protocol.ArtifactChunkUploadResponse) error {
	if !response.Accepted {
		return nil
	}
	entries, err := os.ReadDir(queue.dir)
	if err != nil {
		return fmt.Errorf("read artifact queue: %w", err)
	}
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}
		path := filepath.Join(queue.dir, entry.Name())
		chunk, err := readArtifactChunk(path)
		if err != nil {
			return err
		}
		if chunk.TransferID == response.TransferID && chunk.ArtifactID == response.ArtifactID && chunk.ChunkIndex == response.ChunkIndex {
			if err := os.Remove(path); err != nil {
				return fmt.Errorf("remove acknowledged artifact queue chunk: %w", err)
			}
			if err := os.Remove(artifactChunkPayloadPath(path)); err != nil && !os.IsNotExist(err) {
				return fmt.Errorf("remove acknowledged artifact queue payload: %w", err)
			}
		}
	}
	return nil
}

func (queue ArtifactQueue) chunkPath(chunk protocol.ArtifactChunkUploadRequest) string {
	transferID := sanitizeSegmentName(chunk.TransferID)
	artifactID := sanitizeSegmentName(chunk.ArtifactID)
	return filepath.Join(queue.dir, fmt.Sprintf("%s-%s-%020d.json", transferID, artifactID, chunk.ChunkIndex))
}

func artifactChunkPayloadPath(metadataPath string) string {
	return strings.TrimSuffix(metadataPath, ".json") + ".bin"
}

func writeArtifactQueueFile(path string, payload []byte) error {
	tmp := path + ".tmp"
	file, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open artifact queue file: %w", err)
	}
	_, writeErr := file.Write(payload)
	closeErr := file.Close()
	if writeErr != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("write artifact queue file: %w", writeErr)
	}
	if closeErr != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("close artifact queue file: %w", closeErr)
	}
	if err := syncFile(tmp); err != nil {
		_ = os.Remove(tmp)
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("commit artifact queue file: %w", err)
	}
	return nil
}