diff --git a/README.md b/README.md index f190457f..f84ba3c3 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,8 @@ For production, TAG ships manifests, Compose files, and guides under [`deploy/`] - **Parquet workloads** — `cache.parquet_optimization` warms footer blocks and `cache.meta_on_write` caches object metadata before readers need them; 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/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) +} diff --git a/docs/deploy.md b/docs/deploy.md index 2be457d2..091be263 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -76,9 +76,11 @@ images: ### Graceful shutdown -On SIGINT or SIGTERM, TAG drains admitted HTTP handlers for up to 30 seconds. If the drain expires, it cancels active request contexts and closes their HTTP connections, then keeps the embedded cache open until each synchronous handler and its cache-consistency cleanup returns. A request that finishes inside the drain window retains its normal response; a longer request may be aborted. Cleanup that does not return can keep TAG running beyond the HTTP drain window. +On SIGINT or SIGTERM, TAG stops admitting HTTP requests and drains admitted handlers for up to 30 seconds. A request that finishes inside the drain window retains its normal response. When the drain deadline passes with handlers still running, TAG cancels their request contexts and closes their HTTP connections, logs a warning with the in-flight count, and then keeps waiting for those handlers and their cache-consistency cleanup to return; a longer request may be aborted, but its handler is still joined before shutdown proceeds. -The base StatefulSet sets `terminationGracePeriodSeconds: 30`, which is a separate process-termination limit. Kubernetes may force-stop the pod when that grace expires even if TAG is still waiting for handler cleanup. Set the pod's termination grace to the longest cleanup interval you need to allow; no finite grace can guarantee completion for uncooperative storage. +After the handlers return, TAG 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, and TAG logs a warning with the outstanding task count if the join passes 30 seconds. An admitted tiered cleanup continues under its own timeout so it may finish its conditional S3 delete and metadata repair; the join still waits for it. + +The base StatefulSet sets `terminationGracePeriodSeconds: 30`, which is a separate process-termination limit. When that grace expires, Kubernetes sends SIGKILL even if TAG is still waiting for handler cleanup or detached cache work. Set the pod's termination grace to the longest cleanup interval you need to allow; no finite grace can guarantee completion for uncooperative storage. ### High Availability diff --git a/proxy/block_cache.go b/proxy/block_cache.go index 0f24ebe9..7bc1cc26 100644 --- a/proxy/block_cache.go +++ b/proxy/block_cache.go @@ -1107,8 +1107,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 @@ -1132,7 +1132,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) @@ -1400,7 +1400,12 @@ func (s *Service) beginBlockFetchWithContext(ctx context.Context, blockKey, buck 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 } @@ -1564,8 +1569,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() @@ -1671,22 +1676,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 @@ -1852,17 +1855,17 @@ 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() s.populateBlockMode(ctx, bucket, key, accessKey, secretKey, meta, touched) - }() + }) } // populateBlockMode captures the decision-time token before fetching and commits the block-mode // metadata only after every touched block is durable. Keeping the operation body synchronous makes // its token-failure and mid-fetch invalidation outcomes directly testable at the publication owner; -// triggerBlockModePopulate still runs it under the same detached, bounded context. +// triggerBlockModePopulate still runs it under the same service-owned, bounded context. func (s *Service) populateBlockMode(ctx context.Context, bucket, key, accessKey, secretKey string, meta *cache.CachedObjectMeta, touched []int64) { // Decision-time token BEFORE fetching, so an invalidation that lands // mid-populate bumps the fence past it and the meta commit loses diff --git a/proxy/block_cache_parquet.go b/proxy/block_cache_parquet.go index defbb7bf..cb76af42 100644 --- a/proxy/block_cache_parquet.go +++ b/proxy/block_cache_parquet.go @@ -109,21 +109,25 @@ func (s *Service) maybePrefetchParquetFooter(bucket, key, accessKey, secretKey s s.releaseParquetFooterPrefetch() return } - go func() { - // Delete the marker before freeing capacity so active read markers can - // never temporarily exceed the admission limit. - defer func() { - s.activeBackgroundFetches.Delete(dedupKey) - s.releaseParquetFooterPrefetch() - }() + // Delete the marker before freeing capacity so active read markers can + // never temporarily exceed the admission limit. A trigger that Stop refuses + // to admit never ran, so it returns both exactly as a finished scan would. + finish := func() { + s.activeBackgroundFetches.Delete(dedupKey) + s.releaseParquetFooterPrefetch() + } + if !s.startBackgroundTask(nil, func(ctx context.Context) { + defer finish() // Cooldown only after a completed scan or a stable negative trailer result. // Recording it unconditionally would apply the full window to a shed or // transient failure, turning one miss under load into minutes of silence -- // 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{}{}) } - }() + }) { + finish() + } } // tryAcquireParquetFooterPrefetch admits one read-triggered scan without @@ -164,7 +168,11 @@ func (s *Service) parquetFooterPrefetchWanted(key string, meta *cache.CachedObje // It reports whether the scan completed or found a stable negative trailer result. // A transient cache read or populate failure remains retryable and must not start cooldown. 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, complete := s.parquetFooterLength(ctx, bucket, key, meta, trailer) @@ -319,16 +327,18 @@ func (s *Service) warmParquetFooterOnWriteWithETag(r *http.Request, bucket, key, 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, writtenETag) - }() + s.warmParquetFooterBlocks(ctx, bucket, key, accessKey, secretKey, writtenETag) + }) { + 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, writtenETag string) { - ctx, cancel := context.WithTimeout(context.Background(), backgroundFetchTimeout) +func (s *Service) warmParquetFooterBlocks(parent context.Context, bucket, key, accessKey, secretKey, writtenETag string) { + ctx, cancel := context.WithTimeout(parent, backgroundFetchTimeout) defer cancel() // Stamp BEFORE the upstream trailer read, not after: an invalidation landing diff --git a/proxy/block_cache_parquet_admission_test.go b/proxy/block_cache_parquet_admission_test.go index a9c47f87..bb88edd5 100644 --- a/proxy/block_cache_parquet_admission_test.go +++ b/proxy/block_cache_parquet_admission_test.go @@ -456,3 +456,45 @@ func TestWarmParquetFooterOnWrite_IgnoresReadPrefetchLimit(t *testing.T) { t.Fatalf("write warm after read-limit saturation = (found=%t, err=%v), want metadata", found, err) } } + +func TestMaybePrefetchParquetFooter_RefusedAfterStopReleasesAdmissionSlot(t *testing.T) { + f := newParquetAdmissionFixture(t) + f.service.parquetFooterPrefetchSlots = make(chan struct{}, 1) + t.Cleanup(f.release) + meta := f.seedObject(t, "a.parquet") + + // Nothing is admitted yet, so Stop seals admission and returns at once. A + // trigger that arrives afterwards has already taken its permit and marker. + f.service.Stop() + shedBefore := testutil.ToFloat64(metrics.CacheParquetFooterPrefetchShed) + f.service.maybePrefetchParquetFooter( + "b", "a.parquet", "access", "secret", meta, + meta.ContentLength-parquetTrailerSize, meta.ContentLength-1, nil, + ) + + if got := len(f.service.parquetFooterPrefetchSlots); got != 0 { + t.Fatalf("read-prefetch slots held after a refused trigger = %d, want 0", got) + } + if !f.service.tryAcquireParquetFooterPrefetch() { + t.Fatal("refused trigger leaked the only read-prefetch slot") + } + f.service.releaseParquetFooterPrefetch() + if got := countBackgroundFetchesWithPrefix(f.service, "pq:"); got != 0 { + t.Fatalf("refused trigger left %d read-prefetch markers, want 0", got) + } + if got := len(f.started); got != 0 { + t.Fatalf("refused trigger started %d remote probes, want 0", got) + } + if _, recent := f.service.recentFooterWork.Get("b/a.parquet/\"v1\""); recent { + t.Fatal("refused trigger started a cooldown") + } + if got := testutil.ToFloat64(metrics.CacheParquetFooterPrefetchShed) - shedBefore; got != 0 { + t.Fatalf("stop-time refusal counted as a load shed (delta %v)", got) + } + + // The write-triggered warm shares the marker map and must drop its marker too. + f.service.warmParquetFooterOnWrite(httptest.NewRequest(http.MethodPut, "/b/write.parquet", nil), "b", "write.parquet") + if _, held := f.service.activeBackgroundFetches.Load("pqw:b/write.parquet"); held { + t.Fatal("refused write warm left its marker installed") + } +} 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 c03595e1..de029f09 100644 --- a/proxy/caching.go +++ b/proxy/caching.go @@ -198,8 +198,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 } @@ -210,11 +217,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 @@ -234,7 +244,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()) @@ -256,6 +266,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 @@ -562,7 +582,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 @@ -578,27 +598,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 { @@ -606,11 +638,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 { @@ -660,12 +687,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) @@ -682,7 +708,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 5a99b53d..373de11c 100644 --- a/proxy/get_object.go +++ b/proxy/get_object.go @@ -327,11 +327,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) @@ -361,7 +364,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. @@ -858,14 +861,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..f75d2c6e --- /dev/null +++ b/proxy/lifecycle.go @@ -0,0 +1,152 @@ +package proxy + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "github.com/rs/zerolog/log" +) + +// stopJoinWarnAfter bounds how long Stop joins admitted work silently. The join +// itself has no deadline: it matches main's HTTP drain timeout so a shutdown +// that outlives both is logged once rather than appearing hung. +const stopJoinWarnAfter = 30 * time.Second + +// 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 + active atomic.Int64 // admitted tasks not yet done; reported when Stop's join runs long + stopped chan struct{} // closed once the first stop call has joined every task +} + +func newServiceWorkers() *serviceWorkers { + ctx, cancel := context.WithCancel(context.Background()) + return &serviceWorkers{ctx: ctx, cancel: cancel, stopped: make(chan struct{})} +} + +// 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.active.Add(1) + w.mu.Unlock() + + ctx, cancel := context.WithCancel(parent) + stopServiceCancel := context.AfterFunc(w.ctx, cancel) + done = func() { + stopServiceCancel() + cancel() + w.active.Add(-1) + 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; +// the caller that sealed admission performs the join, so its deadline log is +// emitted once. Later callers wait on a channel rather than a mutex so a +// synctest bubble sees them as durably blocked. +func (w *serviceWorkers) stop() { + w.mu.Lock() + first := !w.stopping + w.stopping = true + w.mu.Unlock() + + if first { + w.cancel() + w.join() + close(w.stopped) + } + <-w.stopped +} + +// join waits for every admitted task without a bound, logging once when the +// wait passes stopJoinWarnAfter and once more when it then completes, so an +// operator can tell a slow cache writer from a hung shutdown. +func (w *serviceWorkers) join() { + joined := make(chan struct{}) + go func() { + w.wg.Wait() + close(joined) + }() + + deadline := time.NewTimer(stopJoinWarnAfter) + defer deadline.Stop() + select { + case <-joined: + return + case <-deadline.C: + } + joinStart := w.warnJoinPastDeadline() + <-joined + log.Info().Dur("waited", time.Since(joinStart)).Msg("Detached cache work finished after the join deadline") +} + +// warnJoinPastDeadline logs once when admitted tasks are still running past +// stopJoinWarnAfter, since stop then waits for them without a bound. It returns +// when the unbounded wait began. +func (w *serviceWorkers) warnJoinPastDeadline() time.Time { + log.Warn(). + Dur("deadline", stopJoinWarnAfter). + Int64("active_tasks", w.active.Load()). + Msg("Service worker join deadline passed; shutdown is waiting past the deadline for detached cache work to finish") + return time.Now() +} + +// 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. A +// join that outlives stopJoinWarnAfter is logged once with the outstanding +// task count, and again when it completes. +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..6be2155a --- /dev/null +++ b/proxy/lifecycle_test.go @@ -0,0 +1,299 @@ +package proxy + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" + "testing" + "testing/synctest" + "time" + + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" + 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 +} + +// lockedLogBuffer collects zerolog output written from the Stop goroutine so the +// test can read it after synctest.Wait without a data race. +type lockedLogBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *lockedLogBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *lockedLogBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +func TestServiceWorkersStopWarnsOnceWhenJoinOutlivesDeadline(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + const ( + warnMsg = "Service worker join deadline passed" + doneMsg = "Detached cache work finished after the join deadline" + ) + logs := &lockedLogBuffer{} + previousLogger := log.Logger + log.Logger = zerolog.New(logs) + t.Cleanup(func() { log.Logger = previousLogger }) + + w := newServiceWorkers() + release := make(chan struct{}) + for range 2 { + if !w.goTask(nil, func(ctx context.Context) { + <-ctx.Done() + <-release + }) { + t.Fatal("task was not admitted before stop") + } + } + + // Two concurrent callers, as main and a signal handler might be: the + // join must still be logged exactly once. + stopped := make(chan struct{}, 2) + for range 2 { + go func() { + w.stop() + stopped <- struct{}{} + }() + } + synctest.Wait() + if w.goTask(nil, func(context.Context) {}) { + t.Fatal("task admitted after stop began") + } + + time.Sleep(stopJoinWarnAfter - time.Second) + synctest.Wait() + if out := logs.String(); strings.Contains(out, warnMsg) { + t.Fatalf("warned before the join deadline:\n%s", out) + } + + time.Sleep(time.Second) + synctest.Wait() + out := logs.String() + if got := strings.Count(out, warnMsg); got != 1 { + t.Fatalf("join-deadline warnings = %d, want 1:\n%s", got, out) + } + if !strings.Contains(out, `"active_tasks":2`) { + t.Fatalf("warning does not report the two outstanding tasks:\n%s", out) + } + if strings.Contains(out, doneMsg) { + t.Fatalf("completion logged while tasks are still held:\n%s", out) + } + select { + case <-stopped: + t.Fatal("stop returned while admitted tasks were still running") + default: + } + + // The held tasks finish a minute past the warning; the completion log + // reports that wait, not the whole join. + time.Sleep(time.Minute) + close(release) + <-stopped + <-stopped + out = logs.String() + if got := strings.Count(out, warnMsg); got != 1 { + t.Fatalf("join-deadline warnings after completion = %d, want 1:\n%s", got, out) + } + if got := strings.Count(out, doneMsg); got != 1 { + t.Fatalf("completion logs = %d, want 1:\n%s", got, out) + } + if !strings.Contains(out, `"waited":60000`) { + t.Fatalf("completion log does not report the 1m wait past the deadline:\n%s", out) + } + }) +} + +func TestServiceWorkersStopIsSilentWhenJoinCompletesInTime(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + logs := &lockedLogBuffer{} + previousLogger := log.Logger + log.Logger = zerolog.New(logs) + t.Cleanup(func() { log.Logger = previousLogger }) + + w := newServiceWorkers() + if !w.goTask(nil, func(ctx context.Context) { <-ctx.Done() }) { + t.Fatal("task was not admitted before stop") + } + w.stop() + time.Sleep(2 * stopJoinWarnAfter) + synctest.Wait() + if out := logs.String(); out != "" { + t.Fatalf("stop that joined in time logged:\n%s", out) + } + }) +} 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 d3ab9e03..6715e8a3 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 4bbc6060..b27aa130 100644 --- a/proxy/service.go +++ b/proxy/service.go @@ -93,6 +93,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 @@ -217,6 +218,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, @@ -1167,7 +1169,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 } @@ -1180,7 +1182,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 } @@ -1191,8 +1193,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 732a7461..ccf723c8 100644 --- a/proxy/tiered.go +++ b/proxy/tiered.go @@ -399,8 +399,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 { @@ -425,7 +425,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 @@ -544,9 +544,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 } @@ -555,15 +560,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() @@ -660,7 +667,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") @@ -859,8 +866,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 @@ -914,6 +925,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 @@ -939,5 +952,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 fb000cb6..5589fa3a 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" @@ -1257,6 +1258,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() }