diff --git a/apps/druid/adapters/cli/callback.go b/apps/druid/adapters/cli/callback.go index cc15448..269cd4a 100644 --- a/apps/druid/adapters/cli/callback.go +++ b/apps/druid/adapters/cli/callback.go @@ -11,6 +11,23 @@ type runtimeCallbackHandler struct { callbacks *appservices.WorkerCallbackManager } +func (h runtimeCallbackHandler) ReportProgress(c *fiber.Ctx) error { + var report struct { + Token string `json:"token"` + Percentage *int64 `json:"percentage"` + } + if err := c.BodyParser(&report); err != nil || report.Percentage == nil { + return fiber.NewError(fiber.StatusBadRequest, "invalid progress report") + } + if *report.Percentage < 0 || *report.Percentage > 100 { + return fiber.NewError(fiber.StatusBadRequest, "percentage must be between 0 and 100") + } + if err := h.callbacks.ReportProgress(c.Params("runtime_id"), report.Token, *report.Percentage); err != nil { + return fiber.NewError(fiber.StatusUnauthorized, err.Error()) + } + return c.SendStatus(fiber.StatusNoContent) +} + func (h runtimeCallbackHandler) CompleteWorker(c *fiber.Ctx, runtimeID callbackapi.Runtime) error { var result callbackapi.WorkerResult if err := c.BodyParser(&result); err != nil { diff --git a/apps/druid/adapters/cli/callback_test.go b/apps/druid/adapters/cli/callback_test.go new file mode 100644 index 0000000..259e11b --- /dev/null +++ b/apps/druid/adapters/cli/callback_test.go @@ -0,0 +1,40 @@ +package cli + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gofiber/fiber/v2" + appservices "github.com/highcard-dev/daemon/apps/druid/core/services" +) + +func TestRuntimeCallbackHandlerReportsProgress(t *testing.T) { + callbacks := appservices.NewWorkerCallbackManager() + token, _, err := callbacks.Register("runtime-1") + if err != nil { + t.Fatal(err) + } + handler := runtimeCallbackHandler{callbacks: callbacks} + app := fiber.New() + app.Post("/internal/v1/workers/:runtime_id/progress", handler.ReportProgress) + + request := httptest.NewRequest( + http.MethodPost, + "/internal/v1/workers/runtime-1/progress", + strings.NewReader(fmt.Sprintf(`{"token":%q,"percentage":42}`, token)), + ) + request.Header.Set(fiber.HeaderContentType, fiber.MIMEApplicationJSON) + response, err := app.Test(request) + if err != nil { + t.Fatal(err) + } + if response.StatusCode != http.StatusNoContent { + t.Fatalf("status = %d; want %d", response.StatusCode, http.StatusNoContent) + } + if progress, ok := callbacks.Progress("runtime-1"); !ok || progress != 42 { + t.Fatalf("progress = %v, %v; want 42, true", progress, ok) + } +} diff --git a/apps/druid/adapters/cli/daemon.go b/apps/druid/adapters/cli/daemon.go index 7369d31..622d142 100644 --- a/apps/druid/adapters/cli/daemon.go +++ b/apps/druid/adapters/cli/daemon.go @@ -170,7 +170,7 @@ func runRuntimeDaemon() error { websocketHandler.SetAllowUnauthenticatedPublic(runtimeAllowUnauthenticatedPublic) handlers := runtimehandlers.RouteHandlers{ Server: runtimehandlers.NewRuntimeServer( - runtimehandlers.NewHealthHandler(), + runtimehandlers.NewHealthHandlerWithProgress(callbacks.Progress), scrollHandler, ), Websocket: websocketHandler, @@ -206,7 +206,9 @@ func runRuntimeDaemon() error { if callbackListener != nil { callbackApp = fiber.New(fiber.Config{DisableStartupMessage: true, ErrorHandler: runtimehandlers.ErrorHandler}) callbackApp.Use(runtimehandlers.RequestLogger) - callbackapi.RegisterHandlers(callbackApp, runtimeCallbackHandler{callbacks: callbacks}) + callbackHandler := runtimeCallbackHandler{callbacks: callbacks} + callbackapi.RegisterHandlers(callbackApp, callbackHandler) + callbackApp.Post("/internal/v1/workers/:runtime_id/progress", callbackHandler.ReportProgress) } return listenRuntimeHTTP(managementApp, publicApp, callbackApp, callbackListener, runtime.Store.StateDir()) } diff --git a/apps/druid/adapters/cli/worker_progress_test.go b/apps/druid/adapters/cli/worker_progress_test.go new file mode 100644 index 0000000..42f7818 --- /dev/null +++ b/apps/druid/adapters/cli/worker_progress_test.go @@ -0,0 +1,53 @@ +package cli + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/highcard-dev/daemon/internal/core/domain" + "github.com/highcard-dev/daemon/internal/core/ports" +) + +func TestWorkerProgressReporterReadsSnapshotProgress(t *testing.T) { + reports := make(chan int64, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + var report struct { + Token string `json:"token"` + Percentage int64 `json:"percentage"` + } + if err := json.NewDecoder(request.Body).Decode(&report); err != nil { + t.Error(err) + } + if report.Token != "token" { + t.Errorf("token = %q; want token", report.Token) + } + reports <- report.Percentage + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + progress := domain.NewSnapshotProgress() + progress.Percentage.Store(37) + stop := startWorkerProgressReporter( + ports.RuntimeWorkerAction{ + RuntimeID: "runtime-1", + CallbackURL: server.URL + "/internal/v1/workers/runtime-1/complete", + CallbackToken: "token", + }, + progress, + time.Hour, + ) + defer stop() + + select { + case percentage := <-reports: + if percentage != 37 { + t.Fatalf("percentage = %d; want 37", percentage) + } + case <-time.After(time.Second): + t.Fatal("progress was not reported") + } +} diff --git a/apps/druid/adapters/cli/worker_pull.go b/apps/druid/adapters/cli/worker_pull.go index 1f7e975..1a9ec26 100644 --- a/apps/druid/adapters/cli/worker_pull.go +++ b/apps/druid/adapters/cli/worker_pull.go @@ -1,13 +1,16 @@ package cli import ( + "bytes" "context" "encoding/json" "fmt" "io" + "net/http" "os" "path/filepath" "strings" + "sync" "time" "github.com/highcard-dev/daemon/internal/callbackapi" @@ -65,6 +68,10 @@ func runWorkerPull(action ports.RuntimeWorkerAction) ports.RuntimeWorkerResult { if root == "" { root = "/scroll" } + progress := domain.NewSnapshotProgress() + stopProgress := startWorkerProgressReporter(action, progress, time.Second) + defer stopProgress() + oci := registry.NewOciClient(loadWorkerRegistryStore()) digest, err := oci.ResolveDigest(action.Artifact) if err == nil { @@ -72,11 +79,11 @@ func runWorkerPull(action ports.RuntimeWorkerAction) ports.RuntimeWorkerResult { } switch action.Mode { case ports.RuntimeWorkerModeUpdate: - err = pullWorkerUpdate(root, action.Artifact, oci) + err = pullWorkerUpdate(root, action.Artifact, oci, progress) case ports.RuntimeWorkerModeRestore: - err = pullWorkerRestore(root, action.Artifact, oci) + err = pullWorkerRestore(root, action.Artifact, oci, progress) default: - err = pullWorkerCreate(root, action.Artifact, oci) + err = pullWorkerCreate(root, action.Artifact, oci, progress) } if err != nil { result.Error = err.Error() @@ -115,7 +122,7 @@ func loadWorkerRegistryStore() *registry.CredentialStore { return registry.NewCredentialStore(config.Registries) } -func pullWorkerCreate(root string, artifact string, oci ports.OciRegistryInterface) error { +func pullWorkerCreate(root string, artifact string, oci ports.OciRegistryInterface, progress *domain.SnapshotProgress) error { if err := os.MkdirAll(root, 0755); err != nil { return err } @@ -137,16 +144,16 @@ func pullWorkerCreate(root string, artifact string, oci ports.OciRegistryInterfa } return copyPath(artifact, root) } - return oci.PullSelective(root, artifact, true, nil) + return oci.PullSelective(root, artifact, true, progress) } -func pullWorkerUpdate(root string, artifact string, oci ports.OciRegistryInterface) error { +func pullWorkerUpdate(root string, artifact string, oci ports.OciRegistryInterface, progress *domain.SnapshotProgress) error { tmp, err := os.MkdirTemp("", "druid-worker-update-*") if err != nil { return err } defer os.RemoveAll(tmp) - if err := coreservices.MaterializeScrollArtifact(artifact, tmp, oci, true); err != nil { + if err := coreservices.MaterializeScrollArtifactWithProgress(artifact, tmp, oci, true, progress); err != nil { return err } scrollYAML, err := os.ReadFile(filepath.Join(tmp, "scroll.yaml")) @@ -162,13 +169,13 @@ func pullWorkerUpdate(root string, artifact string, oci ports.OciRegistryInterfa return mergePulledRoot(tmp, root, skipData) } -func pullWorkerRestore(root string, artifact string, oci ports.OciRegistryInterface) error { +func pullWorkerRestore(root string, artifact string, oci ports.OciRegistryInterface, progress *domain.SnapshotProgress) error { tmp, err := os.MkdirTemp("", "druid-worker-restore-*") if err != nil { return err } defer os.RemoveAll(tmp) - if err := coreservices.MaterializeScrollArtifact(artifact, tmp, oci, true); err != nil { + if err := coreservices.MaterializeScrollArtifactWithProgress(artifact, tmp, oci, true, progress); err != nil { return err } if err := os.MkdirAll(root, 0755); err != nil { @@ -313,6 +320,62 @@ func copyPath(src string, dst string) error { return err } +func startWorkerProgressReporter(action ports.RuntimeWorkerAction, progress *domain.SnapshotProgress, interval time.Duration) func() { + if action.CallbackURL == "" || action.CallbackToken == "" || progress == nil { + return func() {} + } + suffix := "/internal/v1/workers/" + action.RuntimeID + "/complete" + baseURL := strings.TrimSuffix(action.CallbackURL, suffix) + if baseURL == action.CallbackURL { + return func() {} + } + progressURL := baseURL + "/internal/v1/workers/" + action.RuntimeID + "/progress" + client := &http.Client{Timeout: 3 * time.Second} + done := make(chan struct{}) + var wait sync.WaitGroup + wait.Add(1) + go func() { + defer wait.Done() + ticker := time.NewTicker(interval) + defer ticker.Stop() + lastPercentage := int64(-1) + report := func() { + percentage := progress.Percentage.Load() + if percentage == lastPercentage { + return + } + body, _ := json.Marshal(struct { + Token string `json:"token"` + Percentage int64 `json:"percentage"` + }{action.CallbackToken, percentage}) + request, _ := http.NewRequest(http.MethodPost, progressURL, bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + response, err := client.Do(request) + if err == nil { + response.Body.Close() + if response.StatusCode < http.StatusBadRequest { + lastPercentage = percentage + } + } + } + report() + for { + select { + case <-ticker.C: + report() + case <-done: + report() + return + } + } + }() + var once sync.Once + return func() { + once.Do(func() { close(done) }) + wait.Wait() + } +} + func reportWorkerResult(action ports.RuntimeWorkerAction, result ports.RuntimeWorkerResult) error { if action.CallbackURL == "" { body, err := json.Marshal(result) diff --git a/apps/druid/adapters/cli/worker_test.go b/apps/druid/adapters/cli/worker_test.go index 302145b..5fe506c 100644 --- a/apps/druid/adapters/cli/worker_test.go +++ b/apps/druid/adapters/cli/worker_test.go @@ -57,7 +57,7 @@ func TestWorkerRestoreStagesBeforeReplacingRoot(t *testing.T) { mustWrite(t, filepath.Join(root, "data", "old-only.txt"), "old") oci := fakeRestoreOCI{t: t} - if err := pullWorkerRestore(root, "registry.local/backup:1", oci); err != nil { + if err := pullWorkerRestore(root, "registry.local/backup:1", oci, nil); err != nil { t.Fatal(err) } diff --git a/apps/druid/adapters/http/handlers/health_handler.go b/apps/druid/adapters/http/handlers/health_handler.go index 0079f5e..e7658d0 100644 --- a/apps/druid/adapters/http/handlers/health_handler.go +++ b/apps/druid/adapters/http/handlers/health_handler.go @@ -5,12 +5,27 @@ import ( "github.com/highcard-dev/daemon/internal/api" ) -type HealthHandler struct{} +type ProgressLookup func(runtimeID string) (float64, bool) + +type HealthHandler struct { + progress ProgressLookup +} func NewHealthHandler() *HealthHandler { return &HealthHandler{} } +func NewHealthHandlerWithProgress(progress ProgressLookup) *HealthHandler { + return &HealthHandler{progress: progress} +} + func (h *HealthHandler) GetHealthAuth(c *fiber.Ctx) error { - return c.JSON(api.HealthResponse{Mode: "ok"}) + health := api.HealthResponse{Mode: "ok"} + if h.progress != nil { + if progress, ok := h.progress(c.Params("id")); ok { + value := float32(progress) + health.Progress = &value + } + } + return c.JSON(health) } diff --git a/apps/druid/adapters/http/handlers/health_handler_test.go b/apps/druid/adapters/http/handlers/health_handler_test.go new file mode 100644 index 0000000..8678e6b --- /dev/null +++ b/apps/druid/adapters/http/handlers/health_handler_test.go @@ -0,0 +1,33 @@ +package handlers + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gofiber/fiber/v2" + "github.com/highcard-dev/daemon/internal/api" +) + +func TestGetHealthAuthIncludesPullProgress(t *testing.T) { + handler := NewHealthHandlerWithProgress(func(runtimeID string) (float64, bool) { + return 37, runtimeID == "scroll-1" + }) + app := fiber.New() + app.Get("/:id/api/v1/health", handler.GetHealthAuth) + + response, err := app.Test(httptest.NewRequest(http.MethodGet, "/scroll-1/api/v1/health", nil)) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + + var health api.HealthResponse + if err := json.NewDecoder(response.Body).Decode(&health); err != nil { + t.Fatal(err) + } + if health.Progress == nil || *health.Progress != 37 { + t.Fatalf("progress = %v; want 37", health.Progress) + } +} diff --git a/apps/druid/core/services/worker_callbacks.go b/apps/druid/core/services/worker_callbacks.go index 69a2700..dba3775 100644 --- a/apps/druid/core/services/worker_callbacks.go +++ b/apps/druid/core/services/worker_callbacks.go @@ -10,8 +10,9 @@ import ( ) type WorkerCallbackManager struct { - mu sync.Mutex - actions map[string]workerCallbackAction + mu sync.Mutex + actions map[string]workerCallbackAction + progress map[string]int64 } type workerCallbackAction struct { @@ -20,7 +21,10 @@ type workerCallbackAction struct { } func NewWorkerCallbackManager() *WorkerCallbackManager { - return &WorkerCallbackManager{actions: map[string]workerCallbackAction{}} + return &WorkerCallbackManager{ + actions: map[string]workerCallbackAction{}, + progress: map[string]int64{}, + } } func (m *WorkerCallbackManager) Register(runtimeID string) (string, <-chan ports.RuntimeWorkerResult, error) { @@ -36,6 +40,7 @@ func (m *WorkerCallbackManager) Register(runtimeID string) (string, <-chan ports return "", nil, fmt.Errorf("worker action already pending for runtime %s", runtimeID) } m.actions[runtimeID] = workerCallbackAction{token: token, result: ch} + m.progress[runtimeID] = 0 m.mu.Unlock() return token, ch, nil } @@ -43,9 +48,31 @@ func (m *WorkerCallbackManager) Register(runtimeID string) (string, <-chan ports func (m *WorkerCallbackManager) Cancel(runtimeID string) { m.mu.Lock() delete(m.actions, runtimeID) + delete(m.progress, runtimeID) m.mu.Unlock() } +func (m *WorkerCallbackManager) ReportProgress(runtimeID string, token string, percentage int64) error { + m.mu.Lock() + defer m.mu.Unlock() + action, ok := m.actions[runtimeID] + if !ok { + return fmt.Errorf("unknown or completed worker action") + } + if token == "" || token != action.token { + return fmt.Errorf("invalid worker token") + } + m.progress[runtimeID] = max(0, min(100, percentage)) + return nil +} + +func (m *WorkerCallbackManager) Progress(runtimeID string) (float64, bool) { + m.mu.Lock() + defer m.mu.Unlock() + progress, ok := m.progress[runtimeID] + return float64(progress), ok +} + func (m *WorkerCallbackManager) Complete(runtimeID string, token string, result ports.RuntimeWorkerResult) error { m.mu.Lock() action, ok := m.actions[runtimeID] @@ -58,6 +85,7 @@ func (m *WorkerCallbackManager) Complete(runtimeID string, token string, result return fmt.Errorf("invalid worker token") } delete(m.actions, runtimeID) + delete(m.progress, runtimeID) m.mu.Unlock() action.result <- result close(action.result) diff --git a/apps/druid/core/services/worker_callbacks_test.go b/apps/druid/core/services/worker_callbacks_test.go index 2511b9e..c9dd517 100644 --- a/apps/druid/core/services/worker_callbacks_test.go +++ b/apps/druid/core/services/worker_callbacks_test.go @@ -53,3 +53,29 @@ func TestWorkerCallbackRejectsUnknownRuntime(t *testing.T) { t.Fatal("unknown runtime should fail") } } + +func TestWorkerCallbackTracksPullProgress(t *testing.T) { + manager := NewWorkerCallbackManager() + token, _, err := manager.Register("scroll-a") + if err != nil { + t.Fatal(err) + } + + if progress, ok := manager.Progress("scroll-a"); !ok || progress != 0 { + t.Fatalf("initial progress = %v, %v; want 0, true", progress, ok) + } + if err := manager.ReportProgress("scroll-a", "wrong-token", 42); err == nil { + t.Fatal("invalid progress token should fail") + } + if err := manager.ReportProgress("scroll-a", token, 42); err != nil { + t.Fatal(err) + } + if progress, ok := manager.Progress("scroll-a"); !ok || progress != 42 { + t.Fatalf("reported progress = %v, %v; want 42, true", progress, ok) + } + + manager.Cancel("scroll-a") + if _, ok := manager.Progress("scroll-a"); ok { + t.Fatal("cancelled progress should be removed") + } +} diff --git a/internal/core/services/registry/oci.go b/internal/core/services/registry/oci.go index 9d9dbdc..226ebfd 100644 --- a/internal/core/services/registry/oci.go +++ b/internal/core/services/registry/oci.go @@ -201,6 +201,13 @@ func (c *OciClient) PullSelective(dir string, artifact string, includeData bool, if progress != nil { progress.Mode.Store(domain.SnapshotProgressModeRestore) progress.Percentage.Store(0) + defer progress.Mode.Store(domain.SnapshotProgressModeIdle) + } + storeProgress := func(done, total int64) { + if progress == nil || total <= 0 { + return + } + progress.Percentage.Store(min(99, done*100/total)) } copyOpts := oras.CopyOptions{ @@ -259,10 +266,7 @@ func (c *OciClient) PullSelective(dir string, artifact string, includeData bool, done := completed.Add(1) total := totalLayers.Load() bytesDownloaded.Add(desc.Size) - if progress != nil && total > 0 { - pct := done * 100 / total - progress.Percentage.Store(pct) - } + storeProgress(done, total) title := desc.Annotations["org.opencontainers.image.title"] logger.Log().Debug("Pulled layer", zap.String("title", title), @@ -279,10 +283,7 @@ func (c *OciClient) PullSelective(dir string, artifact string, includeData bool, done := completed.Add(1) total := totalLayers.Load() bytesDownloaded.Add(desc.Size) - if progress != nil && total > 0 { - pct := done * 100 / total - progress.Percentage.Store(pct) - } + storeProgress(done, total) title := desc.Annotations["org.opencontainers.image.title"] logger.Log().Debug("Layer already exists locally, skipped", zap.String("title", title), @@ -303,17 +304,9 @@ func (c *OciClient) PullSelective(dir string, artifact string, includeData bool, manifestDescriptor, err := oras.Copy(ctx, repoInstance, ref, fs, dstRef, copyOpts) stopProgress() if err != nil { - if progress != nil { - progress.Mode.Store(domain.SnapshotProgressModeIdle) - } return err } - if progress != nil { - progress.Percentage.Store(100) - progress.Mode.Store(domain.SnapshotProgressModeIdle) - } - logger.Log().Info("Manifest pulled", zap.String("digest", manifestDescriptor.Digest.String()), zap.String("mediaType", manifestDescriptor.MediaType)) jsonData, err := json.Marshal(&manifestDescriptor) @@ -349,6 +342,10 @@ func (c *OciClient) PullSelective(dir string, artifact string, includeData bool, return fmt.Errorf("failed to write annotations: %w", err) } + if progress != nil { + progress.Percentage.Store(100) + } + return nil } diff --git a/internal/core/services/registry/oci_test.go b/internal/core/services/registry/oci_test.go index 2179b3d..e8d5a84 100644 --- a/internal/core/services/registry/oci_test.go +++ b/internal/core/services/registry/oci_test.go @@ -117,6 +117,47 @@ func fakeRegistry(t *testing.T) *httptest.Server { return srv } +func TestPullSelectiveDoesNotReportCompleteBeforeMetadataIsWritten(t *testing.T) { + t.Chdir(t.TempDir()) + server := fakeRegistry(t) + registryHost := strings.TrimPrefix(server.URL, "http://") + source := filepath.Join("scrolls", "progress-test") + if err := os.MkdirAll(source, 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile( + filepath.Join(source, "scroll.yaml"), + []byte("name: progress-test\nversion: 0.1.0\napp_version: \"1.0\"\n"), + 0644, + ); err != nil { + t.Fatal(err) + } + + client := &OciClient{ + credentialStore: NewCredentialStore(nil), + plainHTTP: true, + } + repository := registryHost + "/test/progress" + if _, err := client.Push(source, repository, "1.0", nil, false, nil); err != nil { + t.Fatal(err) + } + + destination := filepath.Join("pull", "progress-test") + if err := os.MkdirAll(filepath.Join(destination, "manifest.json"), 0755); err != nil { + t.Fatal(err) + } + progress := domain.NewSnapshotProgress() + if err := client.PullSelective(destination, repository+":1.0", true, progress); err == nil { + t.Fatal("pull should fail when manifest.json cannot be written") + } + if percentage := progress.Percentage.Load(); percentage >= 100 { + t.Fatalf("failed pull progress = %d; want below 100", percentage) + } + if mode := progress.Mode.Load(); mode != domain.SnapshotProgressModeIdle { + t.Fatalf("failed pull mode = %v; want idle", mode) + } +} + func TestValidateCredentialsUsesPlainHTTPEnv(t *testing.T) { t.Setenv("DRUID_REGISTRY_PLAIN_HTTP", "true") diff --git a/internal/core/services/runtime_scroll_manager.go b/internal/core/services/runtime_scroll_manager.go index a431cb1..ce0aba5 100644 --- a/internal/core/services/runtime_scroll_manager.go +++ b/internal/core/services/runtime_scroll_manager.go @@ -102,6 +102,10 @@ func RuntimeScrollIDFromName(name string) string { } func MaterializeScrollArtifact(artifact string, root string, ociRegistry ports.OciRegistryInterface, includeData bool) error { + return MaterializeScrollArtifactWithProgress(artifact, root, ociRegistry, includeData, nil) +} + +func MaterializeScrollArtifactWithProgress(artifact string, root string, ociRegistry ports.OciRegistryInterface, includeData bool, progress *domain.SnapshotProgress) error { if artifact == "" { return fmt.Errorf("artifact is required") } @@ -126,7 +130,7 @@ func MaterializeScrollArtifact(artifact string, root string, ociRegistry ports.O if ociRegistry == nil { return fmt.Errorf("OCI registry is required to pull %s", artifact) } - if err := ociRegistry.PullSelective(root, artifact, includeData, nil); err != nil { + if err := ociRegistry.PullSelective(root, artifact, includeData, progress); err != nil { return err } return os.MkdirAll(filepath.Join(root, domain.RuntimeDataDir), 0755)