From 84ae52583b59e4299e6bf23bb354c59be853f5c3 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Wed, 14 Oct 2026 06:16:08 +0000 Subject: [PATCH 1/3] test: cover shutdown storage ordering Signed-off-by: Perfloop Agent --- cmd/tag/main_shutdown_test.go | 715 ++++++++++++++++++++++++++++++++++ 1 file changed, 715 insertions(+) create mode 100644 cmd/tag/main_shutdown_test.go diff --git a/cmd/tag/main_shutdown_test.go b/cmd/tag/main_shutdown_test.go new file mode 100644 index 00000000..330cb551 --- /dev/null +++ b/cmd/tag/main_shutdown_test.go @@ -0,0 +1,715 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "errors" + "flag" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strconv" + "strings" + "sync" + "sync/atomic" + "syscall" + "testing" + "time" + + "github.com/rs/zerolog/log" + "github.com/tigrisdata/tag/auth" + "github.com/tigrisdata/tag/cache" +) + +const ( + shutdownProbeBucket = "shutdown-probe" + shutdownProbeKey = "range-object" + shutdownProbeETag = `"shutdown-etag"` + + shutdownTestChildEnv = "TAG_SHUTDOWN_TEST_CHILD" + shutdownTestConfigEnv = "TAG_SHUTDOWN_TEST_CONFIG" + shutdownTestHoldEnv = "TAG_SHUTDOWN_TEST_HOLD" +) + +var errShutdownOrderingViolation = errors.New("embedded storage closed before the detached cache writer exited") + +type storageCloseObserver interface { + IsClosed() bool +} + +type shutdownTestChildState struct { + storageMu sync.RWMutex + storage storageCloseObserver + + writerActive atomic.Bool + + releaseWriter chan struct{} + releaseWriterOnce sync.Once + finish chan struct{} + finishOnce sync.Once + signalHandled chan struct{} + signalHandledOnce sync.Once +} + +var shutdownTestChild atomic.Pointer[shutdownTestChildState] + +func (s *shutdownTestChildState) setStorage(storage storageCloseObserver) { + s.storageMu.Lock() + s.storage = storage + s.storageMu.Unlock() +} + +func (s *shutdownTestChildState) storageClosed() bool { + s.storageMu.RLock() + storage := s.storage + s.storageMu.RUnlock() + return storage != nil && storage.IsClosed() +} + +func (s *shutdownTestChildState) event(name string) { + fmt.Fprintf(os.Stdout, "EVENT %s closed=%t writer=%t\n", name, s.storageClosed(), s.writerActive.Load()) +} + +type shutdownSignalLogWriter struct { + state *shutdownTestChildState + once sync.Once +} + +func (w *shutdownSignalLogWriter) Write(p []byte) (int, error) { + _, _ = os.Stderr.Write(p) + if bytes.Contains(p, []byte("Received shutdown signal")) { + w.once.Do(func() { + w.state.signalHandledOnce.Do(func() { close(w.state.signalHandled) }) + w.state.event("shutdown-signal-received") + }) + } + return len(p), nil +} + +// PutStream pauses the real embedded-client operation at its cache boundary. +// After release it delegates to embedded storage unless Close has already run; +// in that case it reports the forbidden ordering without touching a closed DB. +func (c *embeddedBlockCacheClient) PutStream(ctx context.Context, key string, r io.Reader, ttlSeconds int64) error { + state := shutdownTestChild.Load() + if os.Getenv(shutdownTestChildEnv) == "1" && state != nil { + state.setStorage(c.Client.Storage()) + } + if os.Getenv(shutdownTestChildEnv) != "1" || os.Getenv(shutdownTestHoldEnv) != "true" || key != cache.MakeBodyKey(shutdownProbeBucket, shutdownProbeKey, shutdownProbeETag) { + return c.Client.PutStream(ctx, key, r, ttlSeconds) + } + + if state == nil { + return errors.New("shutdown probe state is unavailable") + } + // This operation is reached only after the detached writer's service-task admission. + state.event("worker-admitted") + state.writerActive.Store(true) + state.event("writer-entered") + select { + case <-ctx.Done(): + if errors.Is(ctx.Err(), context.Canceled) { + state.event("writer-canceled") + } else { + state.event("writer-deadline") + } + <-state.releaseWriter + case <-state.releaseWriter: + } + + if c.Client.Storage().IsClosed() { + state.event("cache-after-close") + state.writerActive.Store(false) + state.event("writer-exited") + return errors.New("embedded storage closed while a cache write was active") + } + err := c.Client.PutStream(ctx, key, r, ttlSeconds) + state.writerActive.Store(false) + state.event("writer-exited") + return err +} + +// TestShutdownChild is the subprocess entry for the real-main shutdown probe. +func TestShutdownChild(t *testing.T) { + if os.Getenv(shutdownTestChildEnv) != "1" { + t.Skip("subprocess helper") + } + + state := &shutdownTestChildState{ + releaseWriter: make(chan struct{}), + finish: make(chan struct{}), + signalHandled: make(chan struct{}), + } + shutdownTestChild.Store(state) + go func() { + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + switch strings.TrimSpace(scanner.Text()) { + case "release": + state.releaseWriterOnce.Do(func() { close(state.releaseWriter) }) + case "finish": + state.finishOnce.Do(func() { close(state.finish) }) + } + } + }() + + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGTERM) + relayStop := make(chan struct{}) + go func() { + select { + case <-signals: + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-state.signalHandled: + return + case <-relayStop: + return + case <-ticker.C: + // The test process catches a signal sent before main has installed + // its listener, then relays it until main reaches its shutdown wait. + _ = syscall.Kill(os.Getpid(), syscall.SIGTERM) + } + } + case <-relayStop: + } + }() + previousLogger := log.Logger + log.Logger = log.Logger.Output(&shutdownSignalLogWriter{state: state}) + + flag.CommandLine = flag.NewFlagSet("tag", flag.ContinueOnError) + os.Args = []string{"tag", "--config", os.Getenv(shutdownTestConfigEnv)} + main() + log.Logger = previousLogger + signal.Stop(signals) + close(relayStop) + state.event("backend-closed") + <-state.finish +} + +func TestShutdownWaitsForBackgroundWriterBeforeEmbeddedClose(t *testing.T) { + if err := runShutdownOrderingProbe(t); err != nil { + t.Fatal(err) + } +} + +func TestShutdownDrainsNormalRequestWithinGrace(t *testing.T) { + if err := runShutdownGraceProbe(); err != nil { + t.Fatal(err) + } +} + +func runShutdownGraceProbe() error { + started := make(chan struct{}, 1) + releaseBody := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseBody) }) } + defer release() + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"grace-etag"`) + w.Header().Set("Content-Length", "8") + w.Header().Set("Content-Type", "application/octet-stream") + w.WriteHeader(http.StatusOK) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + select { + case started <- struct{}{}: + default: + } + select { + case <-releaseBody: + _, _ = io.WriteString(w, "ABCDEFGH") + case <-r.Context().Done(): + } + })) + defer origin.Close() + + child, err := startShutdownChild(origin.URL, false) + if err != nil { + return fmt.Errorf("start TAG subprocess: %w", err) + } + defer child.cleanup() + if err := child.waitReady(); err != nil { + return fmt.Errorf("wait for TAG readiness: %w", err) + } + + tagURL := fmt.Sprintf("http://127.0.0.1:%d", child.port) + signer := auth.NewRequestSigner(tagURL, "us-east-1") + req, err := signer.SignRequest(context.Background(), http.MethodGet, "/grace-bucket/grace-key", nil, "", shutdownTestAccessKey, shutdownTestSecretKey, make(http.Header)) + if err != nil { + return fmt.Errorf("sign in-flight GET: %w", err) + } + type response struct { + status int + body string + err error + } + responseDone := make(chan response, 1) + go func() { + resp, err := (&http.Client{Timeout: shutdownProbeTimeout}).Do(req) + if err != nil { + responseDone <- response{err: err} + return + } + body, readErr := io.ReadAll(resp.Body) + closeErr := resp.Body.Close() + responseDone <- response{status: resp.StatusCode, body: string(body), err: errors.Join(readErr, closeErr)} + }() + select { + case <-started: + case <-time.After(shutdownProbeTimeout): + release() + return errors.New("normal GET did not reach the controlled origin") + } + if err := child.signal(syscall.SIGTERM); err != nil { + release() + return fmt.Errorf("send SIGTERM during normal request: %w", err) + } + if _, err := child.waitLine("EVENT shutdown-signal-received ", shutdownProbeTimeout); err != nil { + release() + return fmt.Errorf("main did not receive SIGTERM during normal request: %w", err) + } + release() + + select { + case result := <-responseDone: + if result.err != nil { + return fmt.Errorf("normal request during HTTP drain: %w", result.err) + } + if result.status != http.StatusOK || result.body != "ABCDEFGH" { + return fmt.Errorf("normal request during HTTP drain = (%d, %q), want (200, ABCDEFGH)", result.status, result.body) + } + case <-time.After(shutdownProbeTimeout): + return errors.New("normal request did not finish within the HTTP grace") + } + closed, err := child.waitLine("EVENT backend-closed ", shutdownProbeTimeout) + if err != nil { + return fmt.Errorf("wait for embedded storage close: %w", err) + } + if !strings.Contains(closed, "closed=true") { + return fmt.Errorf("embedded storage did not report closed after graceful drain: %s", closed) + } + child.finish() + if err := child.waitExit(shutdownProbeTimeout); err != nil { + return fmt.Errorf("TAG subprocess exit: %w", err) + } + if err := checkShutdownDrainResult(child); err != nil { + return err + } + return nil +} + +func runShutdownOrderingProbe(t *testing.T) error { + t.Helper() + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", shutdownProbeETag) + w.Header().Set("Content-Type", "application/octet-stream") + if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", "bytes 0-3/8") + w.Header().Set("Content-Length", "4") + w.WriteHeader(http.StatusPartialContent) + _, _ = io.WriteString(w, "ABCD") + return + } + w.Header().Set("Content-Length", "8") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "ABCDEFGH") + })) + defer origin.Close() + + child, err := startShutdownChild(origin.URL, true) + if err != nil { + return fmt.Errorf("start TAG subprocess: %w", err) + } + defer func() { + child.cleanup() + child.recordEvents(t) + }() + if err := child.waitReady(); err != nil { + return fmt.Errorf("wait for TAG readiness: %w", err) + } + + tagURL := fmt.Sprintf("http://127.0.0.1:%d", child.port) + headers := make(http.Header) + headers.Set("Range", "bytes=0-3") + signer := auth.NewRequestSigner(tagURL, "us-east-1") + req, err := signer.SignRequest(context.Background(), http.MethodGet, "/"+shutdownProbeBucket+"/"+shutdownProbeKey, nil, "", shutdownTestAccessKey, shutdownTestSecretKey, headers) + if err != nil { + return fmt.Errorf("sign cold range request: %w", err) + } + resp, err := (&http.Client{Timeout: shutdownProbeTimeout}).Do(req) + if err != nil { + return fmt.Errorf("cold range request: %w", err) + } + body, readErr := io.ReadAll(resp.Body) + closeErr := resp.Body.Close() + if readErr != nil || closeErr != nil { + return fmt.Errorf("read cold range response: read=%v close=%v", readErr, closeErr) + } + if resp.StatusCode != http.StatusPartialContent || string(body) != "ABCD" { + return fmt.Errorf("cold range response = (%d, %q), want (206, ABCD)", resp.StatusCode, body) + } + if _, err := child.waitLine("EVENT writer-entered ", shutdownProbeTimeout); err != nil { + return fmt.Errorf("background populate did not reach embedded cache PutStream: %w", err) + } + if !child.seen("EVENT worker-admitted ") { + return errors.New("detached cache writer reached PutStream without a recorded admission event") + } + if err := child.signal(syscall.SIGTERM); err != nil { + return fmt.Errorf("send SIGTERM: %w", err) + } + if _, err := child.waitLine("EVENT shutdown-signal-received ", shutdownProbeTimeout); err != nil { + return fmt.Errorf("main did not receive SIGTERM: %w", err) + } + + first, err := child.waitAny([]string{"EVENT writer-canceled ", "EVENT backend-closed "}, shutdownProbeTimeout) + if err != nil { + return fmt.Errorf("wait for shutdown event: %w", err) + } + if strings.HasPrefix(first, "EVENT backend-closed ") { + if err := finishShutdownChild(child); err != nil { + return err + } + if strings.Contains(first, "closed=true") && strings.Contains(first, "writer=true") { + return fmt.Errorf("%w: %s", errShutdownOrderingViolation, first) + } + return fmt.Errorf("shutdown closed storage before a service-cancellation event, but the held writer state was not observable: %s", first) + } + + if err := child.send("release"); err != nil { + return fmt.Errorf("release embedded cache writer: %w", err) + } + order, err := child.waitAny([]string{"EVENT writer-exited ", "EVENT backend-closed "}, shutdownProbeTimeout) + if err != nil { + return fmt.Errorf("wait for writer exit or embedded-storage close: %w", err) + } + if strings.HasPrefix(order, "EVENT backend-closed ") && strings.Contains(order, "writer=true") { + if err := finishShutdownChild(child); err != nil { + return err + } + return fmt.Errorf("%w: %s", errShutdownOrderingViolation, order) + } + if !strings.HasPrefix(order, "EVENT writer-exited ") { + if err := finishShutdownChild(child); err != nil { + return err + } + return fmt.Errorf("unexpected shutdown event before writer exit: %s", order) + } + if child.seen("EVENT cache-after-close ") { + if err := finishShutdownChild(child); err != nil { + return err + } + return fmt.Errorf("%w: cache access was attempted after embedded storage closed", errShutdownOrderingViolation) + } + closed, err := child.waitLine("EVENT backend-closed ", shutdownProbeTimeout) + if err != nil { + return fmt.Errorf("wait for embedded-storage close: %w", err) + } + if !strings.Contains(closed, "closed=true") || !strings.Contains(closed, "writer=false") { + if err := finishShutdownChild(child); err != nil { + return err + } + return fmt.Errorf("backend close did not follow writer exit: %s", closed) + } + return finishShutdownChild(child) +} + +func finishShutdownChild(child *shutdownChildProcess) error { + child.release() + child.finish() + if err := child.waitExit(shutdownProbeTimeout); err != nil { + return fmt.Errorf("TAG subprocess exit: %w", err) + } + return checkShutdownDrainResult(child) +} + +func checkShutdownDrainResult(child *shutdownChildProcess) error { + if output := child.stderr.String(); strings.Contains(output, "Error during server shutdown") { + return fmt.Errorf("HTTP drain did not complete successfully: %s", output) + } + return nil +} + +const ( + shutdownTestAccessKey = "AKIAIOSFODNN7EXAMPLE" + shutdownTestSecretKey = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + shutdownProbeTimeout = 20 * time.Second +) + +type shutdownLockedBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *shutdownLockedBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *shutdownLockedBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +type shutdownChildProcess struct { + command *exec.Cmd + input io.WriteCloser + lines chan string + done chan struct{} + + mu sync.Mutex + waitErr error + history []string + stderr shutdownLockedBuffer + port int + dir string +} + +func startShutdownChild(upstream string, holdWriter bool) (*shutdownChildProcess, error) { + port, err := availableTCPPort() + if err != nil { + return nil, err + } + clusterPort, err := availableTCPPort() + if err != nil { + return nil, err + } + grpcPort, err := availableTCPPort() + if err != nil { + return nil, err + } + + dir, err := os.MkdirTemp("", "tag-shutdown-") + if err != nil { + return nil, err + } + cacheDir := filepath.Join(dir, "cache") + if err := os.Mkdir(cacheDir, 0o700); err != nil { + os.RemoveAll(dir) + return nil, err + } + configPath := filepath.Join(dir, "config.yaml") + configText := fmt.Sprintf("mode: signing\nserver:\n bind_ip: 127.0.0.1\n http_port: %d\nupstream:\n endpoint: %s\n region: us-east-1\ncache:\n enabled: true\n disk_path: %s\n node_id: shutdown-probe\n cluster_addr: 127.0.0.1:%d\n grpc_addr: 127.0.0.1:%d\n grpc_auth: false\n legacy_coordination: false\n block_caching_enabled: false\n size_threshold: 1048576\nlog:\n level: info\n", port, strconv.Quote(upstream), strconv.Quote(cacheDir), clusterPort, grpcPort) + if err := os.WriteFile(configPath, []byte(configText), 0o600); err != nil { + os.RemoveAll(dir) + return nil, err + } + + executable, err := os.Executable() + if err != nil { + os.RemoveAll(dir) + return nil, err + } + command := exec.Command(executable, "-test.run=^TestShutdownChild$") + command.Env = shutdownChildEnvironment(configPath, holdWriter) + input, err := command.StdinPipe() + if err != nil { + os.RemoveAll(dir) + return nil, err + } + stdoutReader, stdoutWriter := io.Pipe() + command.Stdout = stdoutWriter + child := &shutdownChildProcess{ + command: command, + input: input, + lines: make(chan string, 256), + done: make(chan struct{}), + port: port, + dir: dir, + } + command.Stderr = &child.stderr + if err := command.Start(); err != nil { + _ = stdoutReader.Close() + _ = stdoutWriter.Close() + _ = input.Close() + os.RemoveAll(dir) + return nil, err + } + go func() { + scanner := bufio.NewScanner(stdoutReader) + for scanner.Scan() { + child.lines <- scanner.Text() + } + close(child.lines) + _ = stdoutReader.Close() + }() + go func() { + err := command.Wait() + _ = stdoutWriter.Close() + child.mu.Lock() + child.waitErr = err + child.mu.Unlock() + close(child.done) + }() + return child, nil +} + +func availableTCPPort() (int, error) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return 0, err + } + port := listener.Addr().(*net.TCPAddr).Port + if err := listener.Close(); err != nil { + return 0, err + } + return port, nil +} + +func shutdownChildEnvironment(configPath string, holdWriter bool) []string { + env := make([]string, 0, len(os.Environ())+8) + for _, item := range os.Environ() { + name, _, _ := strings.Cut(item, "=") + if strings.HasPrefix(name, "TAG_") || strings.HasPrefix(name, "AWS_") { + continue + } + env = append(env, item) + } + env = append(env, + "AWS_ACCESS_KEY_ID="+shutdownTestAccessKey, + "AWS_SECRET_ACCESS_KEY="+shutdownTestSecretKey, + "AWS_EC2_METADATA_DISABLED=true", + "TAG_LOG_LEVEL=info", + shutdownTestChildEnv+"=1", + shutdownTestConfigEnv+"="+configPath, + shutdownTestHoldEnv+"="+strconv.FormatBool(holdWriter), + ) + return env +} + +func (c *shutdownChildProcess) waitReady() error { + client := &http.Client{Timeout: time.Second} + deadline := time.NewTimer(shutdownProbeTimeout) + defer deadline.Stop() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + url := fmt.Sprintf("http://127.0.0.1:%d/health", c.port) + for { + resp, err := client.Get(url) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + } + select { + case <-deadline.C: + return fmt.Errorf("TAG health endpoint did not become ready: %v", err) + case <-ticker.C: + case <-c.done: + return fmt.Errorf("TAG exited before becoming ready: %w; stderr=%s", c.waitError(), c.stderr.String()) + } + } +} + +func (c *shutdownChildProcess) waitLine(prefix string, timeout time.Duration) (string, error) { + return c.waitAny([]string{prefix}, timeout) +} + +func (c *shutdownChildProcess) waitAny(prefixes []string, timeout time.Duration) (string, error) { + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + select { + case line, ok := <-c.lines: + if !ok { + return "", fmt.Errorf("TAG subprocess output closed: %w", c.waitError()) + } + c.history = append(c.history, line) + for _, prefix := range prefixes { + if strings.HasPrefix(line, prefix) { + return line, nil + } + } + case <-c.done: + return "", fmt.Errorf("TAG subprocess exited: %w", c.waitError()) + case <-timer.C: + return "", fmt.Errorf("timed out waiting for %q; stderr=%s", prefixes, c.stderr.String()) + } + } +} + +func (c *shutdownChildProcess) recordEvents(t *testing.T) { + t.Helper() + for line := range c.lines { + c.history = append(c.history, line) + } + for _, line := range c.history { + if strings.HasPrefix(line, "EVENT ") || strings.HasPrefix(line, "STATUS ") { + t.Log(line) + } + } +} + +func (c *shutdownChildProcess) seen(prefix string) bool { + for _, line := range c.history { + if strings.HasPrefix(line, prefix) { + return true + } + } + return false +} + +func (c *shutdownChildProcess) send(command string) error { + _, err := fmt.Fprintln(c.input, command) + return err +} + +func (c *shutdownChildProcess) release() { + _ = c.send("release") +} + +func (c *shutdownChildProcess) finish() { + _ = c.send("finish") +} + +func (c *shutdownChildProcess) signal(sig os.Signal) error { + if c.command.Process == nil { + return errors.New("TAG subprocess has not started") + } + return c.command.Process.Signal(sig) +} + +func (c *shutdownChildProcess) waitExit(timeout time.Duration) error { + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-c.done: + return c.waitError() + case <-timer.C: + return fmt.Errorf("timed out waiting for TAG subprocess exit: %s", c.stderr.String()) + } +} + +func (c *shutdownChildProcess) waitError() error { + c.mu.Lock() + defer c.mu.Unlock() + return c.waitErr +} + +func (c *shutdownChildProcess) cleanup() { + c.release() + c.finish() + if c.command.Process != nil { + select { + case <-c.done: + default: + _ = c.command.Process.Kill() + } + } + _ = c.waitExit(shutdownProbeTimeout) + _ = c.input.Close() + _ = os.RemoveAll(c.dir) +} From 57d0bf513bd3136d400f68252eae9fdd36f433bc Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Wed, 14 Oct 2026 06:16:08 +0000 Subject: [PATCH 2/3] fix: join cache workers before embedded close Signed-off-by: Perfloop Agent --- README.md | 2 + cmd/tag/main.go | 5 + docs/deploy.md | 4 + proxy/block_cache.go | 51 +++++----- proxy/block_cache_parquet.go | 26 +++-- proxy/caching.go | 84 ++++++++++------ proxy/get_object.go | 27 ++++-- proxy/lifecycle.go | 96 ++++++++++++++++++ proxy/lifecycle_test.go | 177 ++++++++++++++++++++++++++++++++++ proxy/meta_on_write.go | 14 ++- proxy/revalidation.go | 24 ++++- proxy/service.go | 10 +- proxy/tiered.go | 31 ++++-- proxy/tiered_test.go | 76 +++++++++++++++ proxy/write_through.go | 21 ++-- tests/integration/testutil.go | 3 + 16 files changed, 546 insertions(+), 105 deletions(-) create mode 100644 proxy/lifecycle.go create mode 100644 proxy/lifecycle_test.go diff --git a/README.md b/README.md index 4d85e42d..38f6af4f 100644 --- a/README.md +++ b/README.md @@ -127,6 +127,8 @@ For production, TAG ships manifests, Compose files, and guides under [`deploy/`] - **Parquet workloads** — caching parquet metadata ahead of the reader that needs it; see [docs/parquet-optimization.md](docs/parquet-optimization.md). - **Benchmarks** — see [docs/benchmarks.md](docs/benchmarks.md). +On SIGINT or SIGTERM, TAG drains HTTP requests before joining detached cache work and closing embedded storage. A cache operation that does not return after cancellation can extend shutdown; see [Graceful shutdown](docs/deploy.md#graceful-shutdown). + ## Architecture ```text diff --git a/cmd/tag/main.go b/cmd/tag/main.go index 66b0d15d..17f9e331 100644 --- a/cmd/tag/main.go +++ b/cmd/tag/main.go @@ -399,5 +399,10 @@ func main() { log.Error().Err(err).Msg("Error during server shutdown") } + // HTTP draining stops new requests and waits for handlers, but detached + // cache work may still be using embedded storage. Seal and join that work + // before the deferred embeddedCache.Close runs. + service.Stop() + log.Info().Msg("TAG shutdown complete") } diff --git a/docs/deploy.md b/docs/deploy.md index b58f6ea0..3f204c70 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -214,6 +214,10 @@ What to watch: Because segments must pass the age gate first, reclaim begins roughly two hours after a pod restart, not immediately. +### Graceful shutdown + +On SIGINT or SIGTERM, TAG gives HTTP request draining its existing 30-second timeout. After a successful drain, it cancels the Service lifetime and joins admitted detached cache work before closing embedded storage. This join has no deadline: a cache call that does not return after cancellation can extend process shutdown. An admitted tiered cleanup continues under its own timeout so it may finish its conditional S3 delete and metadata repair; Stop still waits for the task. The bundled StatefulSet uses `terminationGracePeriodSeconds: 30`, so Kubernetes can terminate TAG before an uncooperative operation returns; adjust that grace for your workload, but no fixed value guarantees completion. This ordering does not cover handlers that survive the HTTP drain deadline. + ### Health Checks TAG exposes a health endpoint: diff --git a/proxy/block_cache.go b/proxy/block_cache.go index 0581d568..06a5b0d4 100644 --- a/proxy/block_cache.go +++ b/proxy/block_cache.go @@ -839,8 +839,8 @@ func (s *Service) serveFullObjectFromBlockCache( // full GET probes again. if !meta.BlocksComplete && promoTokenOK && s.remainingMetaTTL(meta) > 0 { expectETag := meta.ETag - go func() { - pctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout) + s.startBackgroundTask(nil, func(taskCtx context.Context) { + pctx, cancel := context.WithTimeout(taskCtx, cacheWriteTimeout) defer cancel() // Re-read for the CONTENT to promote (a promoted copy of a stale // snapshot must never clobber a fresh entry) — but commit under @@ -864,7 +864,7 @@ func (s *Service) serveFullObjectFromBlockCache( if _, perr := s.cache.PutMetaIfVersion(pctx, bucket, key, &promoted, remaining, promoToken); perr != nil { log.Debug().Err(perr).Str("bucket", bucket).Str("key", key).Msg("Blocks-complete promotion failed") } - }() + }) } meta.WriteHeaders(w) @@ -1092,7 +1092,12 @@ func (s *Service) beginBlockFetch(blockKey, bucket, key, accessKey, secretKey st s.blockFetches[blockKey] = state s.blockFetchMu.Unlock() - go s.runBlockFetch(state, blockKey, bucket, key, accessKey, secretKey, meta, blockIdx, knownMissing) + if !s.startBackgroundTask(nil, func(ctx context.Context) { + s.runBlockFetch(ctx, state, blockKey, bucket, key, accessKey, secretKey, meta, blockIdx, knownMissing) + }) { + state.completeServe(nil, nil, context.Canceled) + s.finishBlockFetch(blockKey, state, context.Canceled) + } return state } @@ -1253,8 +1258,8 @@ func (s *Service) fetchBlocksForAssembly(ctx context.Context, bucket, key, acces // buffer to a remote writer only after io.ReadFull has completed. Local owners // (and clients that cannot report ownership) retain the prior synchronous write // behavior; remote owners publish the immutable bytes to assemblers first. -func (s *Service) runBlockFetch(state *blockFetchState, blockKey, bucket, key, accessKey, secretKey string, meta *cache.CachedObjectMeta, blockIdx int64, knownMissing bool) { - fetchCtx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) +func (s *Service) runBlockFetch(parent context.Context, state *blockFetchState, blockKey, bucket, key, accessKey, secretKey string, meta *cache.CachedObjectMeta, blockIdx int64, knownMissing bool) { + fetchCtx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) fail := func(err error) { state.completeServe(nil, nil, err) cancel() @@ -1360,22 +1365,20 @@ func (s *Service) runBlockFetch(state *blockFetchState, blockKey, bucket, key, a } // Remote persistence is intentionally detached from the assembled response. - // The state retains both the slot and buffer until PutBlock has consumed the - // bytes, so this cannot create unbounded writers or recycle a live pool buffer. + // This runBlockFetch task remains its owner after publishing the validated + // bytes, and releases the slot and buffer only after PutBlock actually exits. state.completeServe(bufp, blockBuf, nil) - go func() { - perr := s.cache.PutBlock(fetchCtx, bucket, key, meta.ETag, meta.BlockSize, blockIdx, blockBuf, ttl) - if perr != nil { - metrics.CacheBlockPopulateFailed.Inc() - log.Debug().Err(perr).Str("bucket", bucket).Str("key", key).Int64("block", blockIdx).Msg("Detached remote block cache write failed") - } else { - metrics.CacheBlockPopulated.Inc() - metrics.CacheBlockBytesPopulated.Add(float64(blockLen)) - } - s.releaseCacheSlot(reserveWeight) - cancel() - s.finishBlockFetch(blockKey, state, perr) - }() + perr := s.cache.PutBlock(fetchCtx, bucket, key, meta.ETag, meta.BlockSize, blockIdx, blockBuf, ttl) + if perr != nil { + metrics.CacheBlockPopulateFailed.Inc() + log.Debug().Err(perr).Str("bucket", bucket).Str("key", key).Int64("block", blockIdx).Msg("Detached remote block cache write failed") + } else { + metrics.CacheBlockPopulated.Inc() + metrics.CacheBlockBytesPopulated.Add(float64(blockLen)) + } + s.releaseCacheSlot(reserveWeight) + cancel() + s.finishBlockFetch(blockKey, state, perr) } // invalidateStaleMeta deletes the object's cache entry only if the stored meta @@ -1538,8 +1541,8 @@ func (s *Service) triggerBlockModePopulate(bucket, key, accessKey, secretKey str log.Debug().Str("bucket", bucket).Str("key", key).Int("touched", len(touched)).Msg("Block-mode populate skipped - range touches too many blocks (would amplify)") return } - go func() { - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) + s.startBackgroundTask(nil, func(parent context.Context) { + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() // Decision-time token BEFORE fetching, so an invalidation that lands @@ -1558,7 +1561,7 @@ func (s *Service) triggerBlockModePopulate(bucket, key, accessKey, secretKey str return } s.finalizeBlockModeMeta(ctx, bucket, key, meta, len(touched), expected) - }() + }) } // finalizeBlockModeMeta stamps the block-mode meta once its blocks have landed — the diff --git a/proxy/block_cache_parquet.go b/proxy/block_cache_parquet.go index e6274484..c484d4a4 100644 --- a/proxy/block_cache_parquet.go +++ b/proxy/block_cache_parquet.go @@ -100,16 +100,18 @@ func (s *Service) maybePrefetchParquetFooter(bucket, key, accessKey, secretKey s if _, loaded := s.activeBackgroundFetches.LoadOrStore(dedupKey, struct{}{}); loaded { return } - go func() { + if !s.startBackgroundTask(nil, func(ctx context.Context) { defer s.activeBackgroundFetches.Delete(dedupKey) // Cooldown ONLY on a completed scan. Recording it unconditionally would apply // the full window to a budget shed or a transient fetch failure, turning one // shed under load into minutes of silence and more serial footer misses -- // precisely when the prefetch is most worth retrying. - if s.prefetchParquetFooterBlocks(bucket, key, accessKey, secretKey, meta, trailer) && s.recentFooterWork != nil { + if s.prefetchParquetFooterBlocksWithContext(ctx, bucket, key, accessKey, secretKey, meta, trailer) && s.recentFooterWork != nil { s.recentFooterWork.Add(versionKey, struct{}{}) } - }() + }) { + s.activeBackgroundFetches.Delete(dedupKey) + } } // parquetFooterPrefetchWanted reports whether a completed serve is the signal @@ -134,7 +136,11 @@ func (s *Service) parquetFooterPrefetchWanted(key string, meta *cache.CachedObje // It reports whether the scan ran to completion; a caller may use that to decide // whether to suppress repeat work, which must not happen after a failure. func (s *Service) prefetchParquetFooterBlocks(bucket, key, accessKey, secretKey string, meta *cache.CachedObjectMeta, trailer []byte) bool { - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) + return s.prefetchParquetFooterBlocksWithContext(s.workers.ctx, bucket, key, accessKey, secretKey, meta, trailer) +} + +func (s *Service) prefetchParquetFooterBlocksWithContext(parent context.Context, bucket, key, accessKey, secretKey string, meta *cache.CachedObjectMeta, trailer []byte) bool { + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() footerLen, ok := s.parquetFooterLength(ctx, bucket, key, meta, trailer) @@ -284,16 +290,18 @@ func (s *Service) warmParquetFooterOnWrite(r *http.Request, bucket, key string) if _, loaded := s.activeBackgroundFetches.LoadOrStore(dedupKey, struct{}{}); loaded { return } - go func() { + if !s.startBackgroundTask(nil, func(ctx context.Context) { defer s.activeBackgroundFetches.Delete(dedupKey) - s.warmParquetFooterBlocks(bucket, key, accessKey, secretKey) - }() + s.warmParquetFooterBlocks(ctx, bucket, key, accessKey, secretKey) + }) { + s.activeBackgroundFetches.Delete(dedupKey) + } } // warmParquetFooterBlocks resolves the object's size, ETag and metadata length with a // single suffix-range read, then populates the blocks the metadata spans. -func (s *Service) warmParquetFooterBlocks(bucket, key, accessKey, secretKey string) { - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) +func (s *Service) warmParquetFooterBlocks(parent context.Context, bucket, key, accessKey, secretKey string) { + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() // Stamp BEFORE the upstream trailer read, not after: an invalidation landing diff --git a/proxy/caching.go b/proxy/caching.go index cd82233f..ad650ae3 100644 --- a/proxy/caching.go +++ b/proxy/caching.go @@ -187,8 +187,15 @@ func (s *Service) setupCacheListener( } } + taskCtx, taskDone, admitted := s.beginBackgroundTask(ctx) + if !admitted { + s.releaseCacheSlot(weight) + return nil, nil + } + listener := broadcaster.Subscribe() if listener == nil { + taskDone() s.releaseCacheSlot(weight) return nil, nil } @@ -199,11 +206,14 @@ func (s *Service) setupCacheListener( // Start goroutine to consume chunks, build metadata, and write to cache go func() { + defer taskDone() defer s.releaseCacheSlot(weight) defer close(errCh) - // Wait for headers to build metadata - statusCode, headers, err := listener.WaitForHeaders(ctx) + // Wait for headers to build metadata. This context includes both the + // request lifetime and the service lifetime; after headers arrive the + // cache write below detaches from the request but not from the service. + statusCode, headers, err := listener.WaitForHeaders(taskCtx) if err != nil { pipeWriter.CloseWithError(err) listener.DrainAndRelease() // Return any buffered pooled chunks @@ -223,7 +233,7 @@ func (s *Service) setupCacheListener( // Use a detached context for cache writes to avoid cancellation when HTTP request completes. // Scale timeout by content length so large objects have enough time to write. - cacheCtx, cacheCancel := context.WithTimeout(context.Background(), cacheWriteTimeoutForSize(meta.ContentLength)) + cacheCtx, cacheCancel := context.WithTimeout(s.workers.ctx, cacheWriteTimeoutForSize(meta.ContentLength)) defer cacheCancel() ttl := int(s.config.Cache.TTL.Seconds()) @@ -526,7 +536,7 @@ func (s *Service) fetchFullObjectToCache( } } - cacheCtx, cacheCancel := context.WithTimeout(context.Background(), cacheWriteTimeoutForSize(meta.ContentLength)) + cacheCtx, cacheCancel := context.WithTimeout(s.workers.ctx, cacheWriteTimeoutForSize(meta.ContentLength)) defer cacheCancel() // Embedded local storage reaches storage.Put through a non-seekable reader and // does not observe the context itself. Close the upstream body when the detached @@ -542,27 +552,39 @@ func (s *Service) fetchFullObjectToCache( // the writer until it actually exits; releasing it while an uncooperative cache // client still holds buffers would break the aggregate memory bound. cacheErrCh := make(chan error, 1) - go func() { - var cacheErr error - // Block-eligible full fetches retain the size-based representation used by the - // foreground miss path: blocks are written first and version-preconditioned metadata is - // published last. Smaller objects use the whole-body stream writer. - // The precondition is the trigger's decision-time token: the absence - // token for the absent-gated re-warms, the revalidation picker's - // verdict for a revalidation re-warm (so a FAILED guarded delete - // cannot leave known-stale state the repopulate refuses to repair), - // and warmToken for the warm-after-write paths. A trigger that cannot - // read a token does not fire. - if blockMode { - meta.BlockSize = s.config.Cache.BlockSize - _, cacheErr = s.putBlocksFromStream(cacheCtx, bucket, key, meta, body, ttl, expected) - } else { - _, cacheErr = s.cache.PutWithMetaStreamIfVersion( - cacheCtx, bucket, key, meta, body, ttl, expected, + if !s.startBackgroundTask(cacheCtx, func(writerCtx context.Context) { + // The count and byte reservation belong to the actual cache writer, not + // its supervisor. It can outlive the supervisor's timeout if the cache + // client does not return when its context is canceled. + cacheErr := func() error { + defer s.releaseCacheSlot(weight) + // Block-eligible full fetches retain the size-based representation used by the + // foreground miss path: blocks are written first and version-preconditioned metadata is + // published last. Smaller objects use the whole-body stream writer. + // The precondition is the trigger's decision-time token: the absence + // token for the absent-gated re-warms, the revalidation picker's + // verdict for a revalidation re-warm (so a FAILED guarded delete + // cannot leave known-stale state the repopulate refuses to repair), + // and warmToken for the warm-after-write paths. A trigger that cannot + // read a token does not fire. + if blockMode { + meta.BlockSize = s.config.Cache.BlockSize + _, err := s.putBlocksFromStream(writerCtx, bucket, key, meta, body, ttl, expected) + return err + } + _, err := s.cache.PutWithMetaStreamIfVersion( + writerCtx, bucket, key, meta, body, ttl, expected, ) - } + return err + }() + // Publish only after the reservation is released, so a completed result + // retains the prior synchronous ownership boundary. cacheErrCh <- cacheErr - }() + }) { + return context.Canceled + } + // From here the registered writer owns the reservation through its actual exit. + slotOwned = false cacheErr, timedOut := waitForCacheWrite(cacheCtx, cacheErrCh) if timedOut { @@ -570,11 +592,6 @@ func (s *Service) fetchFullObjectToCache( // queued when the context case wins, and net/http bodies use Close to unblock Read. _ = resp.Body.Close() log.Warn().Str("bucket", bucket).Str("key", key).Msg("Background cache write timeout") - slotOwned = false - go func() { - <-cacheErrCh - s.releaseCacheSlot(weight) - }() return errors.New("cache write timeout") } if cacheErr != nil { @@ -624,12 +641,11 @@ func (s *Service) triggerBackgroundCacheFetch(bucket, key, accessKey, secretKey metrics.RecordBackgroundFetchTriggered() metrics.ActiveBackgroundFetches.Inc() - - go func() { + if !s.startBackgroundTask(s.workers.ctx, func(parent context.Context) { defer metrics.ActiveBackgroundFetches.Dec() defer s.activeBackgroundFetches.Delete(bcastKey) - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() err := s.fetchFullObjectToCache(ctx, bucket, key, accessKey, secretKey, anonymous, prio, expected) @@ -646,7 +662,11 @@ func (s *Service) triggerBackgroundCacheFetch(bucket, key, accessKey, secretKey log.Debug().Str("bucket", bucket).Str("key", key).Msg("Background cache fetch completed") metrics.RecordBackgroundFetchSucceeded() } - }() + }) { + metrics.ActiveBackgroundFetches.Dec() + s.activeBackgroundFetches.Delete(bcastKey) + return + } } // hasNoCacheDirectives reports whether any Cache-Control field has a directive that diff --git a/proxy/get_object.go b/proxy/get_object.go index ed3c1731..9e7a6e93 100644 --- a/proxy/get_object.go +++ b/proxy/get_object.go @@ -307,11 +307,14 @@ func (s *Service) fetchAndBroadcast( // Subscribe ourselves as the first listener listener := broadcaster.Subscribe() - // Start the upstream fetch in a goroutine - go func() { - err := s.streamFromUpstream(ctx, r, bucket, key, accessKey, secretKey, broadcaster) + // Start the upstream fetch in a service-owned task. It is also bound to + // the request context, so client cancellation retains its current behavior. + if !s.startBackgroundTask(ctx, func(fetchCtx context.Context) { + err := s.streamFromUpstream(fetchCtx, r, bucket, key, accessKey, secretKey, broadcaster) broadcaster.Complete(err) - }() + }) { + broadcaster.Complete(context.Canceled) + } // Receive chunks like any other listener err := s.writeChunksToResponse(ctx, w, listener, xCache) @@ -341,7 +344,7 @@ func (s *Service) fetchAndBroadcast( <-broadcaster.Done() // the fetch goroutine above records the upstream outcome if upErr := broadcaster.Error(); upErr == nil || errors.Is(upErr, context.Canceled) || errors.Is(upErr, context.DeadlineExceeded) { - if _, tok, found, gerr := s.cache.GetMetaWithVersion(context.Background(), bucket, key); gerr == nil && !found { + if _, tok, found, gerr := s.cache.GetMetaWithVersion(s.workers.ctx, bucket, key); gerr == nil && !found { // Absent-gated, carrying the absence TOKEN (ocache v1.13.0): // the warm is ordered against a fenced delete landing after // this look, where a bare put-if-absent would recreate over it. @@ -699,14 +702,18 @@ func (s *Service) serveRangeFromCache( // the caller forward to upstream instead of streaming a truncated 206 that the // client cannot distinguish from a valid short read. pr, pw := io.Pipe() - go func() { - streamErr := s.cache.GetRangeStream(ctx, bucket, key, meta.BodyDiscriminator(), rng.start, rng.end, pw) + if !s.startBackgroundTask(ctx, func(cacheCtx context.Context) { + streamErr := s.cache.GetRangeStream(cacheCtx, bucket, key, meta.BodyDiscriminator(), rng.start, rng.end, pw) if streamErr != nil { pw.CloseWithError(streamErr) } else { pw.Close() } - }() + }) { + _ = pw.CloseWithError(context.Canceled) + _ = pr.CloseWithError(context.Canceled) + return false, context.Canceled + } firstByte := make([]byte, 1) n, readErr := pr.Read(firstByte) @@ -830,14 +837,14 @@ func (s *Service) handleRangeWithBackgroundCache( // collisions → wrong bytes). Existing entries are filled incrementally by the serve // path under their own captured BlockSize instead. defer func() { - if _, found, _ := s.cache.GetMeta(context.Background(), bucket, key); !found { + if _, found, _ := s.cache.GetMeta(s.workers.ctx, bucket, key); !found { s.triggerBlockModePopulate(bucket, key, accessKey, secretKey, meta, touched) } }() } } else if cacheable { defer func() { - if _, tok, found, gerr := s.cache.GetMetaWithVersion(context.Background(), bucket, key); gerr == nil && !found { + if _, tok, found, gerr := s.cache.GetMetaWithVersion(s.workers.ctx, bucket, key); gerr == nil && !found { // Absent-gated, carrying the absence TOKEN (ocache v1.13.0): // the warm is ordered against a fenced delete landing after // this look, where a bare put-if-absent would recreate over it. diff --git a/proxy/lifecycle.go b/proxy/lifecycle.go new file mode 100644 index 00000000..62ab59c6 --- /dev/null +++ b/proxy/lifecycle.go @@ -0,0 +1,96 @@ +package proxy + +import ( + "context" + "sync" +) + +// serviceWorkers owns work that may outlive an HTTP handler. The mutex serializes +// admission with Stop: after stopping is set, no caller can add to wg while Stop +// waits on it. +type serviceWorkers struct { + ctx context.Context + cancel context.CancelFunc + + mu sync.Mutex + stopping bool + wg sync.WaitGroup +} + +func newServiceWorkers() *serviceWorkers { + ctx, cancel := context.WithCancel(context.Background()) + return &serviceWorkers{ctx: ctx, cancel: cancel} +} + +// begin admits one unit of service-owned work. parent may add an operation's +// own cancellation (for example an HTTP request); the service lifetime always +// cancels the returned context as well. The caller must call done exactly once. +func (w *serviceWorkers) begin(parent context.Context) (ctx context.Context, done func(), ok bool) { + if parent == nil { + parent = w.ctx + } + + w.mu.Lock() + if w.stopping { + w.mu.Unlock() + return nil, nil, false + } + w.wg.Add(1) + w.mu.Unlock() + + ctx, cancel := context.WithCancel(parent) + stopServiceCancel := context.AfterFunc(w.ctx, cancel) + done = func() { + stopServiceCancel() + cancel() + w.wg.Done() + } + return ctx, done, true +} + +// goTask admits and starts one service-owned task. Descendant work that can +// outlive this function must be admitted separately; joined children may remain +// under this task's ownership. +func (w *serviceWorkers) goTask(parent context.Context, work func(context.Context)) bool { + ctx, done, ok := w.begin(parent) + if !ok { + return false + } + go func() { + defer done() + work(ctx) + }() + return true +} + +// stop seals admission before canceling and waiting. Every concurrent caller +// observes the same cancellation and waits for the same admitted work to exit. +func (w *serviceWorkers) stop() { + w.mu.Lock() + w.stopping = true + w.mu.Unlock() + + w.cancel() + w.wg.Wait() +} + +// Stop cancels the service lifetime and waits until every admitted task has +// exited. Call it after HTTP request draining and before closing the cache +// client. It has no deadline: closing storage while an uncooperative cache call +// is still active would violate the service's resource-ownership boundary. +func (s *Service) Stop() { + if s.workers != nil { + s.workers.stop() + } +} + +func (s *Service) startBackgroundTask(parent context.Context, work func(context.Context)) bool { + return s.workers != nil && s.workers.goTask(parent, work) +} + +func (s *Service) beginBackgroundTask(parent context.Context) (context.Context, func(), bool) { + if s.workers == nil { + return nil, nil, false + } + return s.workers.begin(parent) +} diff --git a/proxy/lifecycle_test.go b/proxy/lifecycle_test.go new file mode 100644 index 00000000..93cc9c70 --- /dev/null +++ b/proxy/lifecycle_test.go @@ -0,0 +1,177 @@ +package proxy + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "sync" + "testing" + "testing/synctest" + "time" + + cacheclient "github.com/tigrisdata/ocache/client" + "github.com/tigrisdata/tag/cache" + "github.com/tigrisdata/tag/config" +) + +var errServiceWorkerOwnership = errors.New("service worker ownership violated") + +type heldCacheWriteClient struct { + cacheclient.CacheClient + entered chan struct{} + canceled chan struct{} + release chan struct{} + exited chan struct{} +} + +func (c *heldCacheWriteClient) PutStream(ctx context.Context, _ string, _ io.Reader, _ int64) error { + close(c.entered) + <-ctx.Done() + close(c.canceled) + <-c.release + close(c.exited) + return ctx.Err() +} + +func TestStopWaitsForWriterAfterFetchSupervisorTimesOut(t *testing.T) { + if err := runStopWaitsForWriterAfterFetchSupervisorTimesOut(t); err != nil { + t.Fatal(err) + } +} + +func runStopWaitsForWriterAfterFetchSupervisorTimesOut(t *testing.T) error { + var proofErr error + synctest.Test(t, func(t *testing.T) { + const ( + bucket = "lifecycle-bucket" + key = "lifecycle-key" + body = "small cached body" + ) + + cfg := config.NewDefault() + cfg.Cache.SetBlockCachingEnabled(false) + cfg.Cache.MaxConcurrentWrites = 1 + cfg.Cache.MaxPopulateMemoryBytes = 8 << 20 + + client := &heldCacheWriteClient{ + CacheClient: cacheclient.NewMemoryCache(), + entered: make(chan struct{}), + canceled: make(chan struct{}), + release: make(chan struct{}), + exited: make(chan struct{}), + } + store := cache.NewCacheWithClient(client, &cfg.Cache) + forwarder := &mockForwarder{ + doFullObjectFunc: func(context.Context, string, string, string, string) (*http.Response, error) { + return cacheableGetResponse(body, `"lifecycle-etag"`), nil + }, + } + svc := NewService(forwarder, store, cfg) + var releaseOnce sync.Once + releaseWriter := func() { releaseOnce.Do(func() { close(client.release) }) } + defer svc.Stop() + defer releaseWriter() + + svc.triggerBackgroundCacheFetch(bucket, key, "access", "secret", false, priorityReadMiss, 0) + select { + case <-client.entered: + t.Log("worker admission observed at embedded PutStream") + case <-time.After(time.Second): + proofErr = errors.New("background populate did not reach the controlled cache writer") + return + } + + // The direct cache writer has a source-defined 60-second minimum while + // the supervisor waits on that context. Synctest advances virtual time, + // so the supervisor can time out while the held client remains active. + time.Sleep(cacheWriteTimeout + time.Second) + synctest.Wait() + if _, active := svc.activeBackgroundFetches.Load(backgroundFetchKey(bucket, key, 0)); active { + proofErr = errors.New("background-fetch supervisor did not leave after its writer timeout") + return + } + select { + case <-client.canceled: + case <-time.After(time.Second): + proofErr = errors.New("cache writer did not observe its source-defined timeout") + return + } + t.Log("held PutStream observed context cancellation from its timeout") + t.Logf("fetch supervisor timed out after %s while the admitted PutStream remained active", cacheWriteTimeout) + if got := len(svc.cacheSemaphore); got == 0 { + proofErr = fmt.Errorf("%w: cache-populate slot released while writer is held", errServiceWorkerOwnership) + return + } else if got != 1 { + proofErr = fmt.Errorf("cache-populate slots while writer is held = %d, want 1", got) + return + } + svc.populateBudget.mu.Lock() + remaining := svc.populateBudget.remaining + svc.populateBudget.mu.Unlock() + wantRemaining := cfg.Cache.MaxPopulateMemoryBytes - backgroundPopulateWriterBufferBytes() + if remaining > wantRemaining { + proofErr = fmt.Errorf("%w: populate budget reservation was released while writer is held", errServiceWorkerOwnership) + return + } else if remaining < wantRemaining { + proofErr = fmt.Errorf("populate budget remaining = %d, want %d while writer is held", remaining, wantRemaining) + return + } + t.Logf("while PutStream is held: cache slots=%d, populate bytes reserved=%d", len(svc.cacheSemaphore), cfg.Cache.MaxPopulateMemoryBytes-remaining) + + stopDone := make(chan struct{}, 2) + go func() { + svc.Stop() + stopDone <- struct{}{} + }() + go func() { + svc.Stop() + stopDone <- struct{}{} + }() + select { + case <-svc.workers.ctx.Done(): + case <-time.After(time.Second): + proofErr = fmt.Errorf("%w: Stop did not cancel the service lifetime", errServiceWorkerOwnership) + return + } + synctest.Wait() + + if svc.startBackgroundTask(nil, func(context.Context) {}) { + proofErr = fmt.Errorf("%w: service admitted detached work after Stop began", errServiceWorkerOwnership) + return + } + t.Log("Stop canceled the service lifetime and rejected a later task admission") + select { + case <-stopDone: + proofErr = fmt.Errorf("%w: Stop returned while the cache writer still owned its reservation", errServiceWorkerOwnership) + return + default: + } + t.Log("both concurrent Stop calls remain blocked while PutStream is held") + + releaseWriter() + select { + case <-client.exited: + case <-time.After(time.Second): + proofErr = errors.New("held cache writer did not exit after release") + return + } + t.Log("inner PutStream exited after its controlled release") + <-stopDone + <-stopDone + svc.populateBudget.mu.Lock() + remaining = svc.populateBudget.remaining + svc.populateBudget.mu.Unlock() + if remaining != cfg.Cache.MaxPopulateMemoryBytes { + proofErr = fmt.Errorf("%w: populate budget remaining after Stop = %d, want %d", errServiceWorkerOwnership, remaining, cfg.Cache.MaxPopulateMemoryBytes) + return + } + if got := len(svc.cacheSemaphore); got != 0 { + proofErr = fmt.Errorf("%w: cache-populate slots after Stop = %d, want 0", errServiceWorkerOwnership, got) + return + } + t.Logf("both Stop calls returned after writer exit; cache slots=%d, populate bytes reserved=%d", len(svc.cacheSemaphore), cfg.Cache.MaxPopulateMemoryBytes-remaining) + }) + return proofErr +} diff --git a/proxy/meta_on_write.go b/proxy/meta_on_write.go index b04d42cb..8c8991b5 100644 --- a/proxy/meta_on_write.go +++ b/proxy/meta_on_write.go @@ -51,10 +51,12 @@ func (s *Service) cacheBlockMetaOnWrite(r *http.Request, bucket, key, writtenETa if _, loaded := s.activeBackgroundFetches.LoadOrStore(dedupKey, struct{}{}); loaded { return } - go func() { + if !s.startBackgroundTask(nil, func(ctx context.Context) { defer s.activeBackgroundFetches.Delete(dedupKey) - s.establishBlockMetaFromHead(bucket, key, accessKey, secretKey, writtenETag) - }() + s.establishBlockMetaFromHeadWithContext(ctx, bucket, key, accessKey, secretKey, writtenETag) + }) { + s.activeBackgroundFetches.Delete(dedupKey) + } } // establishBlockMetaFromHead HEADs the object and stamps a block-mode entry whose @@ -63,7 +65,11 @@ func (s *Service) cacheBlockMetaOnWrite(r *http.Request, bucket, key, writtenETa // an evicted block takes. What it buys is that the first read resolves metadata from // cache instead of paying an upstream round trip to discover it. func (s *Service) establishBlockMetaFromHead(bucket, key, accessKey, secretKey, writtenETag string) { - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) + s.establishBlockMetaFromHeadWithContext(s.workers.ctx, bucket, key, accessKey, secretKey, writtenETag) +} + +func (s *Service) establishBlockMetaFromHeadWithContext(parent context.Context, bucket, key, accessKey, secretKey, writtenETag string) { + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() // Stamp before the HEAD, so an invalidation landing mid-flight is newer than this diff --git a/proxy/revalidation.go b/proxy/revalidation.go index 245baed6..123ce16d 100644 --- a/proxy/revalidation.go +++ b/proxy/revalidation.go @@ -261,14 +261,17 @@ func (s *Service) handleRevalidation200( // whole-mode entry that grew into a block-eligible object must not be re-stored as one whole // blob. Sub-block objects keep the whole-body write. cacheErrCh := make(chan error, 1) - go func() { + if !s.startBackgroundTask(nil, func(parent context.Context) { + cacheCtx, cacheCancel := context.WithTimeout(parent, cacheWriteTimeoutForSize(newMeta.ContentLength)) + defer cacheCancel() + var cacheErr error if s.isBlockEligibleSize(newMeta.ContentLength) { newMeta.BlockSize = s.config.Cache.BlockSize - _, cacheErr = s.putBlocksFromStream(context.Background(), bucket, key, newMeta, pr, ttl, expected) + _, cacheErr = s.putBlocksFromStream(cacheCtx, bucket, key, newMeta, pr, ttl, expected) } else { _, cacheErr = s.cache.PutWithMetaStreamIfVersion( - context.Background(), bucket, key, newMeta, pr, ttl, expected, + cacheCtx, bucket, key, newMeta, pr, ttl, expected, ) } if cacheErr != nil { @@ -279,7 +282,20 @@ func (s *Service) handleRevalidation200( io.Copy(io.Discard, pr) } cacheErrCh <- cacheErr - }() + }) { + _ = pr.CloseWithError(context.Canceled) + _ = pw.CloseWithError(context.Canceled) + // Service shutdown has already closed worker admission. Keep the already + // committed upstream response independent of cache availability. + n, copyErr := io.Copy(w, resp.Body) + metrics.BytesTransferred.WithLabelValues("out").Add(float64(n)) + status := "success" + if copyErr != nil { + status = "error" + } + metrics.RecordRequest("GetObject", status, metrics.SourceUpstream, time.Since(start).Seconds()) + return copyErr + } // Foreground: stream to client via TeeReader (also writes to pipe for cache) teeReader := io.TeeReader(resp.Body, pw) diff --git a/proxy/service.go b/proxy/service.go index 5f43afe9..00d5cdc6 100644 --- a/proxy/service.go +++ b/proxy/service.go @@ -90,6 +90,7 @@ type Service struct { forwarder RequestForwarder cache *cache.Cache config *config.Config + workers *serviceWorkers // Owns detached work until Stop joins it cacheSemaphore chan struct{} // Count ceiling on concurrent cache-populate ops (nil = unlimited) populateBudget *byteBudget // Byte budget bounding all cache buffering — populate + block-serve staging (nil = unlimited) perPopulateCap int64 // Max bytes a foreground broadcast populate can buffer @@ -205,6 +206,7 @@ func NewService(forwarder RequestForwarder, cache *cache.Cache, cfg *config.Conf forwarder: forwarder, cache: cache, config: cfg, + workers: newServiceWorkers(), cacheSemaphore: cacheSem, enginePutSlots: enginePutSlots, populateBudget: populateBudget, @@ -1152,7 +1154,7 @@ func (s *Service) warmOnWrite(r *http.Request, bucket, key string) { // succeeded, the surviving stale row's live version if it failed — the // warm then repairs exactly that state and loses to anything newer, // including a fenced delete landing mid-fetch. - warmTok, ok := s.warmToken(bucket, key) + warmTok, ok := s.warmToken(s.workers.ctx, bucket, key) if !ok { return } @@ -1165,7 +1167,7 @@ func (s *Service) warmOnWrite(r *http.Request, bucket, key string) { if err != nil || accessKey == "" || secretKey == "" { return } - warmTok, ok := s.warmToken(bucket, key) + warmTok, ok := s.warmToken(s.workers.ctx, bucket, key) if !ok { return } @@ -1176,8 +1178,8 @@ func (s *Service) warmOnWrite(r *http.Request, bucket, key string) { // warmToken reads a warm trigger's decision-time token. ok=false means the // token could not be read and the warm must be SKIPPED: expected=0 is the // legacy unordered put-if-absent and would publish over a fence. -func (s *Service) warmToken(bucket, key string) (uint64, bool) { - _, tok, _, err := s.cache.GetMetaWithVersion(context.Background(), bucket, key) +func (s *Service) warmToken(ctx context.Context, bucket, key string) (uint64, bool) { + _, tok, _, err := s.cache.GetMetaWithVersion(ctx, bucket, key) if err != nil { log.Debug().Err(err).Str("bucket", bucket).Str("key", key).Msg("Warm skipped - decision-time token unavailable") return 0, false diff --git a/proxy/tiered.go b/proxy/tiered.go index 66d54c16..739ba239 100644 --- a/proxy/tiered.go +++ b/proxy/tiered.go @@ -395,8 +395,8 @@ func (s *Service) stampUpstreamMarkerAfterCompletion(bucket, key, etag, accessKe if accessKey == "" || secretKey == "" { return } - go func() { - hctx, hcancel := context.WithTimeout(context.Background(), 10*time.Second) + s.startBackgroundTask(nil, func(taskCtx context.Context) { + hctx, hcancel := context.WithTimeout(taskCtx, 10*time.Second) defer hcancel() full, ok := s.headObjectMeta(hctx, bucket, key, accessKey, secretKey, etag) if !ok { @@ -421,7 +421,7 @@ func (s *Service) stampUpstreamMarkerAfterCompletion(bucket, key, etag, accessKe if _, perr := s.cache.PutMetaIfVersion(hctx, bucket, key, full, ttl, curVersion); perr != nil { log.Debug().Err(perr).Str("bucket", bucket).Str("key", key).Msg("Completion marker upgrade failed; unknown-length marker serves until TTL") } - }() + }) } // headObjectMeta HEADs the object with TAG-signed credentials and builds its @@ -537,9 +537,14 @@ func (s *Service) maybeRetierOnRead(bucket, key, accessKey, secretKey string, ma return } } - ctx, cancel := context.WithTimeout(context.Background(), retierFetchTimeout) + taskCtx, taskDone, admitted := s.beginBackgroundTask(nil) + if !admitted { + return + } + ctx, cancel := context.WithTimeout(taskCtx, retierFetchTimeout) if !s.tryRegisterRetier(bucket, key, cancel) { cancel() + taskDone() return } @@ -548,15 +553,17 @@ func (s *Service) maybeRetierOnRead(bucket, key, accessKey, secretKey string, ma // weight is the fixed stream buffering, never the object — a threshold // sized above the populate budget no longer makes markers unhealable. weight := int64(originlessPutStreamWeight) - if !s.acquireCacheSlot(context.Background(), weight, priorityReadMiss) { + if !s.acquireCacheSlot(ctx, weight, priorityReadMiss) { s.unregisterRetier(bucket, key) cancel() + taskDone() metrics.RecordTieredRetier("shed") return } etag := marker.ETag go func() { + defer taskDone() defer s.unregisterRetier(bucket, key) defer s.releaseCacheSlot(weight) defer cancel() @@ -651,7 +658,7 @@ func (s *Service) maybeRetierOnRead(bucket, key, accessKey, secretKey string, ma // Detached context: the fetch ctx is CANCELED by a racing PUT — // precisely the interleaving whose staged body most needs // reclaiming — so the delete must not ride it. - dctx, dcancel := context.WithTimeout(context.Background(), 10*time.Second) + dctx, dcancel := context.WithTimeout(s.workers.ctx, 10*time.Second) defer dcancel() if derr := s.cache.DeleteBody(dctx, bucket, key, ref); derr != nil { log.Debug().Err(derr).Str("bucket", bucket).Str("key", key).Msg("Re-tier staged body delete failed; orphan ages out by TTL") @@ -850,8 +857,12 @@ func (s *Service) deleteUpstreamObjectAsync(bucket, key, etag, accessKey, secret metrics.RecordTieredCleanupSkipped("unverified_backend") return } - go func() { - ctx, cancel := context.WithTimeout(context.Background(), tieredCleanupTimeout) + s.startBackgroundTask(nil, func(parent context.Context) { + // This conditional upstream DELETE and its cache repair are one cleanup + // task. Once admitted, retain their own timeout: canceling after dispatch + // can leave the remote result unknown and a same-ETag marker unrepaired. + // Stop still waits for this task before embedded storage closes. + ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), tieredCleanupTimeout) defer cancel() // The If-Match guard below cannot protect a LIVE upstream object that // carries the SAME ETag: plain-PUT ETags are content MD5, so a @@ -905,6 +916,8 @@ func (s *Service) deleteUpstreamObjectAsync(bucket, key, etag, accessKey, secret // commit) that landed after this read; the version guard refuses // exactly those. One outcome is recorded per cleanup: deleted, // marker_repaired, or repair_failed. + // This task keeps its own timeout after a successful DELETE so its + // observed-version marker repair is not cut short by Service.Stop. cur, curVersion, found, gerr := s.cache.GetMetaWithVersion(ctx, bucket, key) if gerr != nil { // The repair check could not run: a raced-in same-ETag marker may @@ -930,5 +943,5 @@ func (s *Service) deleteUpstreamObjectAsync(bucket, key, etag, accessKey, secret // keeps its metadata; the cleanup itself still completed. metrics.RecordTieredCleanup(resp.StatusCode, nil) } - }() + }) } diff --git a/proxy/tiered_test.go b/proxy/tiered_test.go index c3d5e8ae..27d731d4 100644 --- a/proxy/tiered_test.go +++ b/proxy/tiered_test.go @@ -7,6 +7,7 @@ import ( "encoding/base64" "encoding/hex" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -791,6 +792,81 @@ func TestTieredCleanupRepairsRacedSameETagMarker(t *testing.T) { } } +// A successful upstream DELETE is an irreversible step. Shutdown may cancel +// the cleanup worker at that point, but its observed-version metadata repair +// still has to finish before Stop releases the embedded cache. +func TestTieredCleanupFinishesMarkerRepairDuringServiceStop(t *testing.T) { + mock, _, _ := tieredMock() + var c *cache.Cache + deleteStarted := make(chan struct{}) + releaseDelete := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseDelete) }) } + mock.doObjectDeleteFunc = func(ctx context.Context, bucket, key, etag, accessKey, secretKey string) (*http.Response, error) { + marker := &cache.CachedObjectMeta{ + Bucket: bucket, Key: key, ETag: etag, + BodyUpstream: true, StatusCode: http.StatusOK, ContentLength: 18, + } + _, token, _, err := c.GetMetaWithVersion(ctx, bucket, key) + if err != nil { + return nil, err + } + if wrote, err := c.PutMetaIfVersion(ctx, bucket, key, marker, 60, token); err != nil || !wrote { + return nil, fmt.Errorf("raced marker re-establishment: wrote=%v err=%v", wrote, err) + } + close(deleteStarted) + <-releaseDelete // Service.Stop cancels the lifetime while this task remains owned + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("in-flight cleanup context was canceled: %w", err) + } + return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody}, nil + } + var svc *Service + svc, c = newTieredTestService(mock, 8) + defer svc.Stop() + defer release() + + if w := tieredDo(t, svc, http.MethodPut, "/b/obj", "way past threshold", nil); w.Code != http.StatusOK { + t.Fatalf("large PUT status = %d", w.Code) + } + if w := tieredDo(t, svc, http.MethodPut, "/b/obj", "tiny", nil); w.Code != http.StatusOK { + t.Fatalf("small PUT status = %d", w.Code) + } + select { + case <-deleteStarted: + case <-time.After(2 * time.Second): + t.Fatal("cleanup DELETE did not reach the controlled upstream operation") + } + + stopDone := make(chan struct{}) + go func() { + svc.Stop() + close(stopDone) + }() + select { + case <-svc.workers.ctx.Done(): + case <-time.After(time.Second): + t.Fatal("Stop did not cancel the service lifetime") + } + select { + case <-stopDone: + t.Fatal("Stop returned while the upstream cleanup was still active") + default: + } + + release() + select { + case <-stopDone: + case <-time.After(2 * time.Second): + t.Fatal("Stop did not wait for the post-delete metadata repair") + } + if _, found, err := c.GetMeta(context.Background(), "b", "obj"); err != nil { + t.Fatalf("read metadata after Stop: %v", err) + } else if found { + t.Fatal("same-ETag marker survived the post-delete repair during Stop") + } +} + // bodyPutRaceClient fires race once, on the first PutStream of a body| key — // the window between an unconditional PUT's decision-token read and its meta // commit (the engine writes body first, meta second). diff --git a/proxy/write_through.go b/proxy/write_through.go index 740e22ec..8c190594 100644 --- a/proxy/write_through.go +++ b/proxy/write_through.go @@ -170,14 +170,14 @@ func (s *Service) writeThroughCache(bucket, key string, ts *teeState) bool { body := ts.buf.buf accessKey, secretKey := ts.accessKey, ts.secretKey weight := ts.weight - go func() { - // Hold the tee's populate reservation for the whole goroutine: it accounts for the + if !s.startBackgroundTask(nil, func(ctx context.Context) { + // Hold the tee's populate reservation for the whole task: it accounts for the // object-sized teed body, which stays live until this returns. Releasing it early // (before a fallback allocates) would let the body coexist unaccounted with the // fallback's buffers and push live populate memory past the budget. defer s.releaseCacheSlot(weight) - outcome := s.cacheTeedBodyFromHead(bucket, key, putETag, body, accessKey, secretKey) + outcome := s.cacheTeedBodyFromHead(ctx, bucket, key, putETag, body, accessKey, secretKey) if outcome != teeFallbackWarm { // teeCached: already cached. teeNotCacheable: the authoritative HEAD says the object // isn't cacheable (Cache-Control no-store/private, or over threshold), so a warm would @@ -193,13 +193,16 @@ func (s *Service) writeThroughCache(bucket, key string, ts *teeState) bool { // the slot, which the warm's (separate) blocking acquire then observes. // Token at trigger time, like warmOnWrite: repairs exactly the state // this write left (absence, or the survivor of a failed invalidation). - warmTok, ok := s.warmToken(bucket, key) + warmTok, ok := s.warmToken(ctx, bucket, key) if !ok { return } metrics.WarmOnWriteTriggered.Inc() s.triggerBackgroundCacheFetch(bucket, key, accessKey, secretKey, false /*anonymous*/, priorityWarmWrite, warmTok) - }() + }) { + s.releaseCacheSlot(weight) + return false + } return true } @@ -217,7 +220,7 @@ const ( // the teed body. It returns teeCached on a successful write, teeNotCacheable when the HEAD // authoritatively shows the object should not be cached (so warming is pointless), and // teeFallbackWarm when it couldn't confirm/write the version (so a read-back warm may help). -func (s *Service) cacheTeedBodyFromHead(bucket, key, putETag string, body []byte, accessKey, secretKey string) teeOutcome { +func (s *Service) cacheTeedBodyFromHead(ctx context.Context, bucket, key, putETag string, body []byte, accessKey, secretKey string) teeOutcome { // Decision-time token BEFORE the HEAD (mirroring the populate paths), so // the whole HEAD-to-write window is ordered: a competing overwrite whose // invalidation lands after this read bumps the fence past this token and @@ -225,14 +228,14 @@ func (s *Service) cacheTeedBodyFromHead(bucket, key, putETag string, body []byte // ETag-consistency check below. Read after this PUT's own post-forward // invalidation (this goroutine runs after it), so our own fence doesn't // block us. - _, expected, _, tokErr := s.cache.GetMetaWithVersion(context.Background(), bucket, key) + _, expected, _, tokErr := s.cache.GetMetaWithVersion(ctx, bucket, key) if tokErr != nil { // No token, no ordered commit: fall back to the warm, which reads its // own token (and skips likewise if the store is still failing). return teeFallbackWarm } - headCtx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) + headCtx, cancel := context.WithTimeout(ctx, backgroundFetchTimeout) defer cancel() // Empty etag/lastModified => a plain (non-conditional) HEAD. @@ -268,7 +271,7 @@ func (s *Service) cacheTeedBodyFromHead(bucket, key, putETag string, body []byte } ttl := int(s.config.Cache.TTL.Seconds()) - cacheCtx, cacheCancel := context.WithTimeout(context.Background(), cacheWriteTimeoutForSize(meta.ContentLength)) + cacheCtx, cacheCancel := context.WithTimeout(ctx, cacheWriteTimeoutForSize(meta.ContentLength)) defer cacheCancel() // Put-if-absent: this tee follows the PUT's own invalidation; a racer that // re-established the entry (a warm) caches the current version and wins. diff --git a/tests/integration/testutil.go b/tests/integration/testutil.go index 9d7fab8d..5d52a104 100644 --- a/tests/integration/testutil.go +++ b/tests/integration/testutil.go @@ -497,6 +497,9 @@ func (e *TestEnvironment) Close() { if e.TAGServer != nil { e.TAGServer.Close() } + if e.Service != nil { + e.Service.Stop() + } if e.UpstreamServer != nil { e.UpstreamServer.Close() } From af7e64bf56a6f30b15c6700ed9475e85d1998b63 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Wed, 14 Oct 2026 08:50:14 +0000 Subject: [PATCH 3/3] fix: unblock canceled cache relay writers Signed-off-by: Perfloop Agent --- proxy/cache_populate_completion_test.go | 101 ++++++++++++++++++++++++ proxy/caching.go | 10 +++ 2 files changed, 111 insertions(+) diff --git a/proxy/cache_populate_completion_test.go b/proxy/cache_populate_completion_test.go index 33752da9..eb0c837d 100644 --- a/proxy/cache_populate_completion_test.go +++ b/proxy/cache_populate_completion_test.go @@ -195,6 +195,107 @@ func TestSetupCacheListener_PublishesWhenHeadersReadyAtRequestCancellation(t *te svc.releaseCacheSlot(1) } +// cancelingPipeStreamCache models a remote cache owner that stops consuming a +// streamed write as soon as its service context is canceled. The retained reader +// lets the test release the relay if Stop fails to close the pipe on that return. +type cancelingPipeStreamCache struct { + cacheclient.CacheClient + readerStarted chan struct{} + reader chan io.Reader + writerExited chan struct{} +} + +func (c *cancelingPipeStreamCache) PutStream(ctx context.Context, _ string, r io.Reader, _ int64) error { + var firstByte [1]byte + n, err := r.Read(firstByte[:]) + if n == 0 { + if err != nil { + return err + } + return io.ErrNoProgress + } + + c.reader <- r + close(c.readerStarted) + <-ctx.Done() + close(c.writerExited) + return ctx.Err() +} + +// TestServiceStopUnblocksFullGetWhenCacheWriterStopsReading covers a full GET +// whose detached cache writer has begun a pipe read but returns on shutdown +// without consuming the rest of the body. The HTTP response must finish normally; +// Stop must then join the abandoned relay rather than waiting on its blocked write. +func TestServiceStopUnblocksFullGetWhenCacheWriterStopsReading(t *testing.T) { + const body = "cache relay has more bytes than its reader consumes" + + cfg := config.NewDefault() + cfg.Cache.SetBlockCachingEnabled(false) + cfg.Broadcast.ChunkSize = 8 + client := &cancelingPipeStreamCache{ + CacheClient: cacheclient.NewMemoryCache(), + readerStarted: make(chan struct{}), + reader: make(chan io.Reader, 1), + writerExited: make(chan struct{}), + } + cacheStore := cache.NewCacheWithClient(client, &cfg.Cache) + forwarder := &mockForwarder{ + doRequestFunc: func(context.Context, *http.Request, string, string) (*http.Response, error) { + return cacheableGetResponse(body, `"pipe-shutdown-etag"`), nil + }, + } + svc := NewService(forwarder, cacheStore, cfg) + recorder := httptest.NewRecorder() + requestDone := make(chan error, 1) + go func() { + requestDone <- svc.HandleGetObject(recorder, httptest.NewRequest(http.MethodGet, "/pipe-bucket/pipe-key", nil)) + }() + + select { + case <-client.readerStarted: + case <-time.After(time.Second): + t.Fatal("full GET cache writer did not start reading the body") + } + select { + case err := <-requestDone: + if err != nil { + t.Fatalf("HandleGetObject: %v", err) + } + case <-time.After(time.Second): + t.Fatal("full GET response did not finish while cache writer was held") + } + if recorder.Code != http.StatusOK || recorder.Body.String() != body { + t.Fatalf("full GET response = (%d, %q), want (200, %q)", recorder.Code, recorder.Body.String(), body) + } + + stopDone := make(chan struct{}) + go func() { + svc.Stop() + close(stopDone) + }() + select { + case <-client.writerExited: + case <-time.After(time.Second): + t.Fatal("cache writer did not return after service cancellation") + } + select { + case <-stopDone: + case <-time.After(time.Second): + // The writer has returned, so no cache operation owns this reader any more. + // Drain it only to unblock a failing implementation and leave no goroutine + // behind after reporting the missing pipe-close handoff. + if _, err := io.Copy(io.Discard, <-client.reader); err != nil { + t.Errorf("drain abandoned test reader: %v", err) + } + select { + case <-stopDone: + case <-time.After(time.Second): + t.Fatal("Service.Stop remained blocked after the abandoned reader was drained") + } + t.Fatal("Service.Stop did not join after the cache writer returned") + } +} + // lateFailingStreamCache waits until the test has observed the client handler // return, then fails the stream write. This forces the cache error to arrive // after broadcaster completion rather than before the upstream EOF path returns. diff --git a/proxy/caching.go b/proxy/caching.go index ad650ae3..f9cd5fed 100644 --- a/proxy/caching.go +++ b/proxy/caching.go @@ -255,6 +255,16 @@ func (s *Service) setupCacheListener( cacheErrCh := make(chan error, 1) go func() { var cacheErr error + // Some cache clients can return before consuming their input (for example, + // when a remote stream is canceled). Close the reader as soon as its owner + // exits so the relay cannot block writing bytes no cache operation will read. + defer func() { + if cacheErr != nil { + _ = pipeReader.CloseWithError(cacheErr) + } else { + _ = pipeReader.Close() + } + }() // Block-eligible objects (>= block_size, block caching on) are stored as fixed-size // blocks regardless of how they were fetched — a full-GET read miss or a warm-on-write // read-back. The whole-vs-block boundary is size, not access pattern (RFC 0001): both