126 lines
3.7 KiB
Go
126 lines
3.7 KiB
Go
package spool
|
|
|
|
import (
|
|
"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)
|
|
tmp := path + ".tmp"
|
|
file, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
|
|
if err != nil {
|
|
return fmt.Errorf("open artifact queue chunk: %w", err)
|
|
}
|
|
encodeErr := json.NewEncoder(file).Encode(chunk)
|
|
closeErr := file.Close()
|
|
if encodeErr != nil {
|
|
_ = os.Remove(tmp)
|
|
return fmt.Errorf("encode artifact queue chunk: %w", encodeErr)
|
|
}
|
|
if closeErr != nil {
|
|
_ = os.Remove(tmp)
|
|
return fmt.Errorf("close artifact queue chunk: %w", closeErr)
|
|
}
|
|
if err := os.Rename(tmp, path); err != nil {
|
|
_ = os.Remove(tmp)
|
|
return fmt.Errorf("commit artifact queue chunk: %w", err)
|
|
}
|
|
return 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 {
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open artifact queue chunk: %w", err)
|
|
}
|
|
var chunk protocol.ArtifactChunkUploadRequest
|
|
decodeErr := json.NewDecoder(file).Decode(&chunk)
|
|
closeErr := file.Close()
|
|
if decodeErr != nil {
|
|
return nil, fmt.Errorf("decode artifact queue chunk: %w", decodeErr)
|
|
}
|
|
if closeErr != nil {
|
|
return nil, fmt.Errorf("close artifact queue chunk: %w", closeErr)
|
|
}
|
|
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())
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
return fmt.Errorf("open artifact queue chunk: %w", err)
|
|
}
|
|
var chunk protocol.ArtifactChunkUploadRequest
|
|
decodeErr := json.NewDecoder(file).Decode(&chunk)
|
|
closeErr := file.Close()
|
|
if decodeErr != nil {
|
|
return fmt.Errorf("decode artifact queue chunk: %w", decodeErr)
|
|
}
|
|
if closeErr != nil {
|
|
return fmt.Errorf("close artifact queue chunk: %w", closeErr)
|
|
}
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
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))
|
|
}
|