From 4025bc579bdb7234c3baf258a222ebe4a26e3a3f Mon Sep 17 00:00:00 2001 From: Rafael Matias Date: Wed, 18 Mar 2026 15:53:56 +0100 Subject: [PATCH] feat: wait for block persistence before stopping container and use 30s stop timeout Poll eth_blockNumber after pre-run steps to confirm the client has finished persisting blocks before stopping/checkpointing the container. This prevents data loss with clients like reth that persist blocks asynchronously. Also sets an explicit 30s graceful shutdown timeout on StopContainer for both Docker and Podman. --- pkg/docker/docker.go | 7 +- pkg/podman/podman.go | 7 +- pkg/runner/rpc.go | 103 ++++++++++++++++++++++++++++++ pkg/runner/strategy_checkpoint.go | 18 ++++++ pkg/runner/strategy_container.go | 18 ++++++ 5 files changed, 149 insertions(+), 4 deletions(-) diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go index ed17f90b..87caae06 100644 --- a/pkg/docker/docker.go +++ b/pkg/docker/docker.go @@ -304,9 +304,12 @@ func (m *manager) StartContainer(ctx context.Context, containerID string) error return nil } -// StopContainer stops a container. +// StopContainer stops a container with a 30-second graceful shutdown timeout. func (m *manager) StopContainer(ctx context.Context, containerID string) error { - if err := m.client.ContainerStop(ctx, containerID, container.StopOptions{}); err != nil { + timeout := 30 + if err := m.client.ContainerStop(ctx, containerID, container.StopOptions{ + Timeout: &timeout, + }); err != nil { return fmt.Errorf("stopping container %s: %w", containerID[:12], err) } diff --git a/pkg/podman/podman.go b/pkg/podman/podman.go index 3ef724b8..dcc8209b 100644 --- a/pkg/podman/podman.go +++ b/pkg/podman/podman.go @@ -298,12 +298,15 @@ func (m *manager) StartContainer(ctx context.Context, containerID string) error return nil } -// StopContainer stops a container. +// StopContainer stops a container with a 30-second graceful shutdown timeout. func (m *manager) StopContainer(ctx context.Context, containerID string) error { conn, cancel := m.connWithCtx(ctx) defer cancel() - if err := containers.Stop(conn, containerID, nil); err != nil { + timeout := uint(30) + if err := containers.Stop(conn, containerID, &containers.StopOptions{ + Timeout: &timeout, + }); err != nil { return fmt.Errorf("stopping container %s: %w", containerID[:12], err) } diff --git a/pkg/runner/rpc.go b/pkg/runner/rpc.go index 1565d116..ecad827d 100644 --- a/pkg/runner/rpc.go +++ b/pkg/runner/rpc.go @@ -78,6 +78,109 @@ func (r *runner) checkRPCHealth(ctx context.Context, url string) (string, bool) return rpcResp.Result, true } +const ( + blockPersistenceInterval = 500 * time.Millisecond + blockPersistenceRetries = 60 +) + +// getBlockNumber fetches the current block number via eth_blockNumber. +func (r *runner) getBlockNumber(ctx context.Context, host string, port int) (uint64, error) { + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + url := fmt.Sprintf("http://%s:%d", host, port) + body := `{"jsonrpc":"2.0","method":"eth_blockNumber","params":[],"id":1}` + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(body)) + if err != nil { + return 0, fmt.Errorf("creating request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return 0, fmt.Errorf("executing request: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + return 0, fmt.Errorf("unexpected status code: %d", resp.StatusCode) + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return 0, fmt.Errorf("reading response: %w", err) + } + + var rpcResp struct { + Result string `json:"result"` + } + + if err := json.Unmarshal(respBody, &rpcResp); err != nil { + return 0, fmt.Errorf("parsing response: %w", err) + } + + blockNum, err := strconv.ParseUint(strings.TrimPrefix(rpcResp.Result, "0x"), 16, 64) + if err != nil { + return 0, fmt.Errorf("parsing block number: %w", err) + } + + return blockNum, nil +} + +// waitForBlockPersistence polls eth_blockNumber until it reaches the expected +// block number. Some clients (e.g. reth) may still be persisting blocks in the +// background after accepting payloads, so this ensures the data is fully +// committed before the container is stopped. +func (r *runner) waitForBlockPersistence( + ctx context.Context, + log logrus.FieldLogger, + host string, + port int, + expectedBlock uint64, +) error { + log = log.WithFields(logrus.Fields{ + "expected_block": expectedBlock, + "max_retries": blockPersistenceRetries, + "interval": blockPersistenceInterval, + }) + log.Info("Waiting for block persistence before stopping container") + + for attempt := 1; attempt <= blockPersistenceRetries; attempt++ { + blockNum, err := r.getBlockNumber(ctx, host, port) + if err != nil { + log.WithFields(logrus.Fields{ + "attempt": attempt, + "error": err.Error(), + }).Warn("Failed to query eth_blockNumber, retrying") + } else if blockNum >= expectedBlock { + log.WithFields(logrus.Fields{ + "block_number": blockNum, + "attempts": attempt, + }).Info("Block persistence confirmed") + + return nil + } else { + log.WithFields(logrus.Fields{ + "block_number": blockNum, + "attempt": attempt, + }).Debug("Block not yet persisted, retrying") + } + + select { + case <-time.After(blockPersistenceInterval): + case <-ctx.Done(): + return fmt.Errorf("context cancelled waiting for block persistence: %w", ctx.Err()) + } + } + + return fmt.Errorf( + "block persistence not confirmed after %d attempts (expected block %d)", + blockPersistenceRetries, expectedBlock, + ) +} + // getLatestBlock fetches the latest block number, hash, and state root from the RPC endpoint. func (r *runner) getLatestBlock(ctx context.Context, host string, port int) (uint64, string, string, error) { ctx, cancel := context.WithTimeout(ctx, 10*time.Second) diff --git a/pkg/runner/strategy_checkpoint.go b/pkg/runner/strategy_checkpoint.go index ca225ad5..8dacedd7 100644 --- a/pkg/runner/strategy_checkpoint.go +++ b/pkg/runner/strategy_checkpoint.go @@ -183,6 +183,20 @@ func (r *runner) runTestsWithCheckpointRestore( log.WithField("steps", n).Info("Pre-run steps completed before checkpoint") } + // Wait for the client to finish persisting blocks before checkpointing. + expectedBlock, blkErr := r.getBlockNumber(ctx, containerIP, spec.RPCPort()) + if blkErr != nil { + log.WithError(blkErr).Warn( + "Failed to get current block number before checkpoint, skipping persistence check", + ) + } else if expectedBlock > 0 { + if err := r.waitForBlockPersistence( + ctx, log, containerIP, spec.RPCPort(), expectedBlock, + ); err != nil { + log.WithError(err).Warn("Block persistence wait failed, proceeding with checkpoint") + } + } + // 3. Decide checkpoint export path: tmpfs (RAM) or disk. // // When checkpoint_tmpfs_threshold is configured and the container's @@ -275,10 +289,14 @@ func (r *runner) runTestsWithCheckpointRestore( waitAfterTCPDrop := r.cfg.FullConfig.GetCheckpointWaitAfterTCPDropConns(params.Instance) + log.Info("Checkpointing container (this will stop the container)") + if err := cpMgr.CheckpointContainer(ctx, containerID, exportPath, waitAfterTCPDrop); err != nil { return nil, fmt.Errorf("checkpointing container: %w", err) } + log.Info("Container checkpointed successfully") + defer func() { _ = os.Remove(exportPath) diff --git a/pkg/runner/strategy_container.go b/pkg/runner/strategy_container.go index 16f4d8e4..05205b66 100644 --- a/pkg/runner/strategy_container.go +++ b/pkg/runner/strategy_container.go @@ -101,11 +101,29 @@ func (r *runner) runTestsWithContainerStrategy( ) } + // Wait for the client to finish persisting blocks before stopping. + expectedBlock, blkErr := r.getBlockNumber(ctx, containerIP, spec.RPCPort()) + if blkErr != nil { + log.WithError(blkErr).Warn( + "Failed to get current block number before stop, skipping persistence check", + ) + } else if expectedBlock > 0 { + if err := r.waitForBlockPersistence( + ctx, log, containerIP, spec.RPCPort(), expectedBlock, + ); err != nil { + log.WithError(err).Warn("Block persistence wait failed, proceeding with stop") + } + } + // Stop the initial container so writes are flushed to disk. + log.Info("Stopping container for ZFS snapshot") + if err := r.containerMgr.StopContainer(ctx, containerID); err != nil { return nil, fmt.Errorf("stopping container for ZFS snapshot: %w", err) } + log.Info("Container stopped for ZFS snapshot") + waitForLogDrain(logDone, logCancel, logDrainTimeout) // Sync to flush any dirty pages before snapshotting.