Reconnect: exponential backoff with jitter + max_reconnect_attempts

Project.toml

@@ -1,6 +1,6 @@
 name = "DatabentoAPI"
 uuid = "feb0aed8-3291-460b-9688-521dcb17a4bf"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Tyler Beason <tbeas12@gmail.com>"]
 
 [deps]

src/live/streaming.jl

@@ -18,21 +18,52 @@
 # Reconnect: on TCP drop, rebuild the Live and re-subscribe every schema with
 # its own `start =` the lowest per-instrument timestamp seen so far (ts_recv
 # for BBO families, ts_event otherwise), bounded to within Databento's 24h
-# replay window. `ErrorMsg` from the gateway is treated as terminal for all
-# schemas on this connection (no reconnect).
+# replay window. Between attempts we sleep `rand() * min(cap, base*2^(n-1))`
+# (AWS-style full-jitter exponential backoff, capped at 60s) and give up
+# after `max_reconnect_attempts` retries (default 10; pass `nothing` for
+# unlimited). `ErrorMsg` from the gateway is still terminal for all schemas
+# on this connection (no reconnect).
 #
 # Head-of-line caveat: the reader task `put!`s to each schema's typed channel
 # in turn. If consumer N's disk write falls behind enough to fill its
 # channel_size buffer, reader stalls the TCP socket for every schema. Default
 # channel_size=10_000 is generally enough headroom; tune up for very bursty
 # schemas paired with a slow archive disk.
 
-const _RECONNECT_WAIT_S        = 1.0
+const _RECONNECT_BASE_S        = 1.0
+const _RECONNECT_CAP_S         = 60.0
 const _STALE_THRESHOLD_S       = 30.0
 const _SHUTDOWN_JOIN_TIMEOUT_S = 15.0
 const _CONNECTION_POLL_S       = 0.5
 const _REPLAY_WINDOW_NS        = Int64(24 * 3600) * 1_000_000_000
 
+# Full-jitter exponential backoff (AWS-recommended): wait drawn uniformly from
+# [0, min(cap, base * 2^(attempt-1))). `attempt` is 1-based — `attempt=1` is
+# the delay before the first reconnect.
+function _reconnect_delay(attempt::Integer;
+                          base::Real = _RECONNECT_BASE_S,
+                          cap::Real  = _RECONNECT_CAP_S)
+    attempt < 1 && return 0.0
+    upper = min(float(cap), float(base) * 2.0^(attempt - 1))
+    return rand() * upper
+end
+
+# Sleep up to `delay` seconds, but wake every `_CONNECTION_POLL_S` to check for
+# shutdown / deadline so a long backoff window doesn't pin `duration_s` or
+# `shutdown_requested[]` responsiveness. Matches the cadence of the main
+# coordination loop.
+function _wait_for_reconnect(delay::Float64, ctxs, deadline::Union{Nothing,Float64})
+    delay > 0.0 || return
+    deadline_at = time() + delay
+    while true
+        _any_shutdown(ctxs)                          && return
+        deadline !== nothing && time() >= deadline   && return
+        remaining = deadline_at - time()
+        remaining > 0 || return
+        sleep(min(_CONNECTION_POLL_S, remaining))
+    end
+end
+
 const _SPARSE_SCHEMAS = (Schema.STATUS, Schema.DEFINITION, Schema.IMBALANCE)
 
 # StatusMsg.action codes that warrant a console alert per Databento's status
@@ -370,6 +401,7 @@ function _run_unified_session(ctxs::Dict{Schema.T,SessionContext};
                               stype_in::SType.T,
                               symbols,
                               reconnect::Bool,
+                              max_reconnect_attempts::Union{Integer,Nothing},
                               deadline::Union{Nothing,Float64},
                               key, gateway, port::Integer,
                               wire_compression::Compression.T,
@@ -462,8 +494,14 @@ function _run_unified_session(ctxs::Dict{Schema.T,SessionContext};
         !reconnect                                       && return
         deadline !== nothing && time() >= deadline       && return
         for ctx in values(ctxs); ctx.stats.reconnects += 1; end
-        @warn "live disconnect — reconnecting" schemas=schemas attempt=first(values(ctxs)).stats.reconnects
-        sleep(_RECONNECT_WAIT_S)
+        attempt = first(values(ctxs)).stats.reconnects
+        if max_reconnect_attempts !== nothing && attempt > max_reconnect_attempts
+            @warn "live reconnect attempts exhausted — giving up" schemas=schemas attempts=attempt limit=max_reconnect_attempts
+            return
+        end
+        delay = _reconnect_delay(attempt)
+        @warn "live disconnect — reconnecting" schemas=schemas attempt=attempt delay_s=round(delay; digits=2)
+        _wait_for_reconnect(delay, ctxs, deadline)
     end
 end
 
@@ -538,6 +576,11 @@ Returns the path to the most recently opened output file.
 The default `compress_level = 1` favours throughput; bench measurements show
 L1 is ~35% faster on the write boundary than L3 in exchange for ~15% larger
 files (see `benchmark/PERF_REPORT.md`). Pass a higher level for archival.
+
+`reconnect = true` (default) re-establishes the TCP session after a drop using
+full-jitter exponential backoff (1s base, 60s cap). After
+`max_reconnect_attempts` retries (default 10) the loop gives up and returns;
+pass `nothing` for unlimited.
 """
 function stream_to_file(; schema::Schema.T,
                         symbols,
@@ -550,6 +593,7 @@ function stream_to_file(; schema::Schema.T,
                         compress_level::Integer = 1,
                         rotate_seconds::Union{Nothing,Real} = nothing,
                         reconnect::Bool = true,
+                        max_reconnect_attempts::Union{Integer,Nothing} = 10,
                         key = nothing,
                         gateway = nothing,
                         port::Integer = DEFAULT_LIVE_PORT,
@@ -577,7 +621,9 @@ function stream_to_file(; schema::Schema.T,
     try
         _run_unified_session(ctxs;
             dataset = dataset, stype_in = stype_in, symbols = symbols,
-            reconnect = reconnect, deadline = deadline,
+            reconnect = reconnect,
+            max_reconnect_attempts = max_reconnect_attempts,
+            deadline = deadline,
             key = key, gateway = gateway, port = port,
             wire_compression = wire_compression,
             heartbeat_interval = heartbeat_interval,
@@ -600,6 +646,10 @@ Capture N Live schemas concurrently, one rotating-DBN file per schema. Returns
 a dict mapping each schema to its most recently opened output path. Press
 `Ctrl-C` to stop early; honours `duration_s` if provided.
 
+`reconnect`/`max_reconnect_attempts` work as in [`stream_to_file`](@ref):
+full-jitter exponential backoff between retries (1s base, 60s cap), default
+cap of 10 attempts, pass `nothing` for unlimited.
+
 Example:
 ```julia
 paths = stream_multi_to_files(
@@ -621,6 +671,7 @@ function stream_multi_to_files(; schemas::AbstractVector,
                                compress_level::Integer = 1,
                                rotate_seconds::Union{Nothing,Real} = nothing,
                                reconnect::Bool = true,
+                               max_reconnect_attempts::Union{Integer,Nothing} = 10,
                                key = nothing,
                                gateway = nothing,
                                port::Integer = DEFAULT_LIVE_PORT,
@@ -659,7 +710,9 @@ function stream_multi_to_files(; schemas::AbstractVector,
     session_task = @async try
         _run_unified_session(contexts;
             dataset = dataset, stype_in = stype_in, symbols = symbols,
-            reconnect = reconnect, deadline = deadline,
+            reconnect = reconnect,
+            max_reconnect_attempts = max_reconnect_attempts,
+            deadline = deadline,
             key = key, gateway = gateway, port = port,
             wire_compression = wire_compression,
             heartbeat_interval = heartbeat_interval,

test/runtests.jl

@@ -22,6 +22,7 @@ try
         include("test_live_reader_typed.jl")
         include("test_live_subscribe.jl")
         include("test_live_streaming.jl")
+        include("test_reconnect_backoff.jl")
         include("test_typed_mode_errors.jl")
     end
 catch e

test/test_reconnect_backoff.jl

@@ -0,0 +1,165 @@
+using Test
+using Sockets
+using DatabentoAPI
+using DatabentoAPI: _reconnect_delay, _RECONNECT_BASE_S, _RECONNECT_CAP_S
+using DatabentoBinaryEncoding
+import DatabentoBinaryEncoding as DBN
+
+const _BACKOFF_TEST_KEY = "db-0123456789abcdef0123456789abcdef"
+
+@testset "reconnect — _reconnect_delay full-jitter bounds" begin
+    # attempt=1 → upper = base
+    for _ in 1:200
+        d = _reconnect_delay(1)
+        @test 0 ≤ d < _RECONNECT_BASE_S
+    end
+    # attempt=2 → upper = 2*base
+    for _ in 1:200
+        d = _reconnect_delay(2)
+        @test 0 ≤ d < 2 * _RECONNECT_BASE_S
+    end
+    # attempt=3 → upper = 4*base
+    for _ in 1:200
+        d = _reconnect_delay(3)
+        @test 0 ≤ d < 4 * _RECONNECT_BASE_S
+    end
+    # Cap kicks in once base * 2^(n-1) > cap.
+    # With base=1, cap=60: 2^(n-1) > 60 ⇔ n ≥ 7
+    for n in 7:20
+        for _ in 1:50
+            d = _reconnect_delay(n)
+            @test 0 ≤ d < _RECONNECT_CAP_S
+        end
+    end
+    # attempt < 1 → 0
+    @test _reconnect_delay(0) == 0.0
+    @test _reconnect_delay(-3) == 0.0
+
+    # Custom base/cap kwargs
+    d_custom = _reconnect_delay(1; base = 5.0, cap = 100.0)
+    @test 0 ≤ d_custom < 5.0
+    d_capped = _reconnect_delay(20; base = 1.0, cap = 2.0)
+    @test 0 ≤ d_capped < 2.0
+end
+
+@testset "reconnect — jitter actually varies between calls" begin
+    # 50 draws at attempt=5 should produce many distinct values; if they're
+    # all identical the RNG isn't being consulted (regression guard).
+    draws = [_reconnect_delay(5) for _ in 1:50]
+    @test length(unique(draws)) > 10
+end
+
+@testset "reconnect — max_reconnect_attempts caps connection attempts" begin
+    # Mock gateway that immediately closes each socket. The client will fail
+    # during the LSG-version read, raise, and fall through to the reconnect
+    # path. With max_reconnect_attempts=2 we expect exactly 3 accepts
+    # (1 initial + 2 reconnects), then the loop exits.
+    accept_count = Threads.Atomic{Int}(0)
+    # Pre-stage more scripts than we expect to use; extras simply never get
+    # accepted and the server gets closed in finally.
+    scripts = [function(sock)
+        Threads.atomic_add!(accept_count, 1)
+        # Close immediately — client's first read (LSG version frame) gets EOF.
+        return
+    end for _ in 1:8]
+    mock = spawn_mock_gateway_sequence(scripts)
+
+    t0 = time()
+    mktempdir() do dir
+        DatabentoAPI.stream_to_file(
+            schema = DBN.Schema.TRADES, symbols = ["AAPL"],
+            dataset = "TEST.MOCK", stype_in = DBN.SType.RAW_SYMBOL,
+            base_dir = dir, compress = false,
+            duration_s = 30,                  # generous; the cap should trip first
+            reconnect = true,
+            max_reconnect_attempts = 2,
+            key = _BACKOFF_TEST_KEY,
+            gateway = "127.0.0.1", port = mock.port,
+            wire_compression = DBN.Compression.NONE,
+            heartbeat_log_interval_s = 5.0,
+        )
+    end
+    elapsed = time() - t0
+
+    # 1 initial + 2 reconnects = 3 accepts.
+    @test accept_count[] == 3
+    # Worst-case delay budget: attempt-1 (≤1s) + attempt-2 (≤2s) = ≤3s plus
+    # a few seconds of handshake / cleanup slack. Anything in single digits
+    # means the cap engaged; without the cap this would run for duration_s=30.
+    @test elapsed < 15.0
+
+    # Server cleanup
+    try; close(mock.server); catch; end
+end
+
+@testset "reconnect — duration_s cuts off mid-backoff sleep" begin
+    # Regression guard for Codex P1: with full-jitter backoff a single retry
+    # can sleep up to 60s. The reconnect loop must remain responsive to the
+    # deadline so `duration_s` is honored even if the deadline is crossed
+    # while a backoff sleep is in flight.
+    #
+    # We force the loop into the high-backoff regime by setting a small base
+    # and a large cap via the existing `_reconnect_delay` defaults, then check
+    # that the call returns close to `duration_s` regardless of how many
+    # attempts have stacked up.
+    accept_count = Threads.Atomic{Int}(0)
+    scripts = [function(sock)
+        Threads.atomic_add!(accept_count, 1)
+        return
+    end for _ in 1:50]
+    mock = spawn_mock_gateway_sequence(scripts)
+
+    duration = 1.5
+    t0 = time()
+    mktempdir() do dir
+        DatabentoAPI.stream_to_file(
+            schema = DBN.Schema.TRADES, symbols = ["AAPL"],
+            dataset = "TEST.MOCK", stype_in = DBN.SType.RAW_SYMBOL,
+            base_dir = dir, compress = false,
+            duration_s = duration,
+            reconnect = true,
+            max_reconnect_attempts = nothing,   # let cap-by-deadline drive exit
+            key = _BACKOFF_TEST_KEY,
+            gateway = "127.0.0.1", port = mock.port,
+            wire_compression = DBN.Compression.NONE,
+            heartbeat_log_interval_s = 5.0,
+        )
+    end
+    elapsed = time() - t0
+
+    # Should be roughly `duration_s` plus a small slack for the last in-flight
+    # sleep slice (≤ _CONNECTION_POLL_S = 0.5s) and handshake overhead. Without
+    # the fix, a single sleep of up to 60s could blow this out.
+    @test elapsed < duration + 3.0
+    try; close(mock.server); catch; end
+end
+
+@testset "reconnect — reconnect=false bypasses cap entirely" begin
+    # When reconnect is disabled, the first accepted connection is the only one
+    # regardless of max_reconnect_attempts. Verifies kwargs don't interact
+    # unexpectedly.
+    accept_count = Threads.Atomic{Int}(0)
+    scripts = [function(sock)
+        Threads.atomic_add!(accept_count, 1)
+        return
+    end for _ in 1:5]
+    mock = spawn_mock_gateway_sequence(scripts)
+
+    mktempdir() do dir
+        DatabentoAPI.stream_to_file(
+            schema = DBN.Schema.TRADES, symbols = ["AAPL"],
+            dataset = "TEST.MOCK", stype_in = DBN.SType.RAW_SYMBOL,
+            base_dir = dir, compress = false,
+            duration_s = 30,
+            reconnect = false,
+            max_reconnect_attempts = 10,      # ignored when reconnect=false
+            key = _BACKOFF_TEST_KEY,
+            gateway = "127.0.0.1", port = mock.port,
+            wire_compression = DBN.Compression.NONE,
+            heartbeat_log_interval_s = 5.0,
+        )
+    end
+
+    @test accept_count[] == 1
+    try; close(mock.server); catch; end
+end