Add timeout to download chunks from node (#941)

Layr-Labs · Dec 5, 2024 · b196817 · b196817
1 parent 8ded1ec
commit b196817
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/node/config.go b/node/config.go
@@ -91,6 +91,7 @@ type Config struct {
 
 	EnableV2                    bool
 	OnchainStateRefreshInterval time.Duration
+	ChunkDownloadTimeout        time.Duration
 
 	PprofHttpPort string
 	EnablePprof   bool
@@ -240,6 +241,7 @@ func NewConfig(ctx *cli.Context) (*Config, error) {
 		BLSRemoteSignerEnabled:         blsRemoteSignerEnabled,
 		EnableV2:                       ctx.GlobalBool(flags.EnableV2Flag.Name),
 		OnchainStateRefreshInterval:    ctx.GlobalDuration(flags.OnchainStateRefreshIntervalFlag.Name),
+		ChunkDownloadTimeout:           ctx.GlobalDuration(flags.ChunkDownloadTimeoutFlag.Name),
 		PprofHttpPort:                  ctx.GlobalString(flags.PprofHttpPort.Name),
 		EnablePprof:                    ctx.GlobalBool(flags.EnablePprof.Name),
 	}, nil

diff --git a/node/flags/flags.go b/node/flags/flags.go
@@ -231,6 +231,13 @@ var (
 		EnvVar:   common.PrefixEnvVar(EnvVarPrefix, "ONCHAIN_STATE_REFRESH_INTERVAL"),
 		Value:    1 * time.Hour,
 	}
+	ChunkDownloadTimeoutFlag = cli.DurationFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "chunk-download-timeout"),
+		Usage:    "The timeout for downloading chunks from the relay (default: 30s)",
+		Required: false,
+		EnvVar:   common.PrefixEnvVar(EnvVarPrefix, "CHUNK_DOWNLOAD_TIMEOUT"),
+		Value:    20 * time.Second,
+	}
 
 	// Test only, DO NOT USE the following flags in production
 
@@ -374,6 +381,7 @@ var optionalFlags = []cli.Flag{
 	BLSSignerCertFileFlag,
 	EnableV2Flag,
 	OnchainStateRefreshIntervalFlag,
+	ChunkDownloadTimeoutFlag,
 	PprofHttpPort,
 	EnablePprof,
 }

diff --git a/node/node_v2.go b/node/node_v2.go
@@ -103,7 +103,9 @@ func (n *Node) DownloadBundles(ctx context.Context, batch *corev2.Batch, operato
 		relayKey := relayKey
 		req := requests[relayKey]
 		pool.Submit(func() {
-			bundles, err := relayClient.GetChunksByRange(ctx, relayKey, req.chunkRequests)
+			ctxTimeout, cancel := context.WithTimeout(ctx, n.Config.ChunkDownloadTimeout)
+			defer cancel()
+			bundles, err := relayClient.GetChunksByRange(ctxTimeout, relayKey, req.chunkRequests)
 			if err != nil {
 				n.Logger.Errorf("failed to get chunks from relays: %v", err)
 				bundleChan <- response{