fix(713): EAGAIN retry helper + slot quarantine on STATE_PUT restore failure

common/hydra-socket-retry.h

@@ -0,0 +1,140 @@
+// hydra#713: shared recv-with-EAGAIN-retry helper for non-blocking sockets.
+//
+// Both the M2 state-stream path (llama_io_read_socket::refill in
+// llama-context.cpp) and the RPC framing path (hydra_recv_all in
+// server-context.cpp) need identical EAGAIN/EWOULDBLOCK handling: poll + retry
+// with a bounded deadline, drain on POLLHUP/POLLERR before declaring EOF, and
+// EINTR retry on every syscall.  A single copy here avoids silent drift
+// between the two call sites.
+//
+// Header-only, POSIX only (`#if !defined(_WIN32)`).
+//
+// timeout_ms is a per-recv-call budget (not a per-transfer budget):
+// each call to hydra_recv_with_retry waits at most timeout_ms for the
+// requested n bytes.  The caller's outer loop (hydra_recv_all, refill)
+// may invoke this repeatedly for large transfers.
+
+#ifndef LLAMA_HYDRA_SOCKET_RETRY_H
+#define LLAMA_HYDRA_SOCKET_RETRY_H
+
+#if !defined(_WIN32)
+
+#include <cerrno>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <algorithm>
+#include <chrono>
+
+#include <poll.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+// Attempt a non-blocking recv with bounded poll-retry on EAGAIN/EWOULDBLOCK.
+//
+// Returns:
+//   >0  — bytes read (1..n).  recv may return a short read (r < n), so the
+//         caller's loop must keep reading until n bytes are in or EOF/error.
+//    0  — clean EOF (peer closed after draining any buffered data)
+//   -1  — hard error or timeout; errno is set:
+//           ETIMEDOUT  — timeout_ms elapsed with no data
+//           ECONNRESET / EPIPE / etc. — peer-level failure
+//           EBADF      — bad fd
+//
+// On POLLHUP or POLLERR the helper performs one final recv to drain any
+// data the peer wrote before closing.  Only if that recv also returns 0
+// (EOF) or a hard error does the function return.  This prevents
+// discarding buffered data when the peer writes its final bytes then
+// closes (a common pattern for large STATE_PUT transfers up to 800 MB).
+inline ssize_t hydra_recv_with_retry(int fd, void * buf, size_t n, int timeout_ms) {
+    char * p = reinterpret_cast<char *>(buf);
+
+    // Fast path: attempt recv immediately — no poll overhead for the common case.
+    ssize_t r = ::recv(fd, p, n, 0);
+    if (r > 0) {
+        return r;
+    }
+    if (r == 0) {
+        return 0;  // clean EOF
+    }
+    // r < 0 — check errno before entering the retry loop.
+    if (errno != EAGAIN && errno != EWOULDBLOCK) {
+        return -1;  // hard error (ECONNRESET, EBADF, …)
+    }
+
+    // EAGAIN: poll + retry loop with a bounded wall-clock deadline.
+    // The deadline is relative to the FIRST EAGAIN, not to the original call,
+    // so the caller's per-call budget is respected.
+    //
+    // Elapsed time is measured with a monotonic clock (steady_clock), not by
+    // subtracting the poll slice each iteration.  Slice-subtraction burns the
+    // whole budget on events that consumed no wall time (EINTR, a spurious
+    // POLLIN that still yields EAGAIN), which would time out a healthy
+    // transfer prematurely.  Charging real elapsed time also makes the EINTR
+    // handling actually correct rather than merely bounded.
+    const auto deadline = std::chrono::steady_clock::now()
+        + std::chrono::milliseconds(static_cast<long long>(timeout_ms > 0 ? timeout_ms : 30000));
+
+    for (;;) {
+        const auto now = std::chrono::steady_clock::now();
+        if (now >= deadline) {
+            errno = ETIMEDOUT;
+            return -1;
+        }
+        const auto rem_ms = std::chrono::duration_cast<std::chrono::milliseconds>(deadline - now).count();
+        const int wait_ms = static_cast<int>(std::min<long long>(rem_ms, 1000LL));
+
+        struct pollfd pfd = { fd, POLLIN, 0 };
+        int pr = ::poll(&pfd, 1, wait_ms);
+        if (pr < 0) {
+            if (errno == EINTR) {
+                // Interrupted before any data: real elapsed time is already
+                // charged against the deadline at the top of the loop.
+                continue;
+            }
+            return -1;  // real poll error (EBADF, EINVAL, …)
+        }
+        if (pr == 0) {
+            // Poll slice expired with no data: loop back — the deadline check
+            // accounts for the time that actually passed.
+            continue;
+        }
+
+        // poll returned > 0: at least one event is ready.
+
+        if (pfd.revents & (POLLHUP | POLLERR)) {
+            // On Linux, POLLIN is often set together with POLLHUP when the
+            // peer wrote final bytes then closed.  We must drain buffered
+            // data before declaring EOF — otherwise up to 800 MB of a
+            // legitimate STATE_PUT transfer is silently discarded.
+            r = ::recv(fd, p, n, 0);
+            if (r > 0) {
+                return r;
+            }
+            if (r == 0) {
+                return 0;  // true EOF after drain
+            }
+            // r < 0: recv error after HUP/ERR — propagate.
+            return -1;
+        }
+
+        // POLLIN ready — attempt the actual recv.
+        r = ::recv(fd, p, n, 0);
+        if (r > 0) {
+            return r;
+        }
+        if (r == 0) {
+            return 0;  // clean EOF
+        }
+        // r < 0
+        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
+            // Still no data (spurious readiness / interrupted): retry within
+            // the deadline.  Real elapsed time is charged at the top of the loop.
+            continue;
+        }
+        return -1;  // hard error (ECONNRESET, etc.)
+    }
+}
+
+#endif // !_WIN32
+#endif // LLAMA_HYDRA_SOCKET_RETRY_H

src/llama-context.cpp

@@ -3011,8 +3011,15 @@ size_t llama_context::state_get_size() {
 // hydra: zero-copy socket streaming (class stays here; C wrapper in llama-hydra.cpp)
 #if !defined(_WIN32)
 #include <sys/socket.h>
+// hydra#713: shared EAGAIN/EWOULDBLOCK retry helper (poll + drain-on-HUP).
+#include "../common/hydra-socket-retry.h"
 // xxh3 for M2 decode-side wire-hash verification (see state_seq_set_data_from_fd)
+// Use header-only inline variant so libllama does not need an external
+// xxhash.c object — server-context.cpp already compiles the non-inline
+// implementation for the server binary, but tests link only libllama.
+#define XXH_INLINE_ALL
 #include "../vendor/xxhash/xxhash.h"
+#undef XXH_INLINE_ALL
 
 class llama_io_write_socket : public llama_io_write_i {
     // hydra#334: chunk size is caller-supplied (see llama_cparams::hydra_state_chunk_size,
@@ -3143,9 +3150,13 @@ class llama_io_read_socket : public llama_io_read_i {
         if (staging_pos < staging_len) {
             return;
         }
-        ssize_t r = ::recv(fd, staging.data(), staging.size(), 0);
+        ssize_t r = hydra_recv_with_retry(fd, staging.data(), staging.size(), 30000);
         if (r <= 0) {
-            throw std::runtime_error("hydra: socket recv failed during state restore");
+            if (r == 0) {
+                throw std::runtime_error("hydra: socket recv EOF during state restore");
+            }
+            throw std::runtime_error(std::string("hydra: socket recv failed during state restore: ")
+                                     + std::strerror(errno));
         }
         staging_pos = 0;
         staging_len = (size_t)r;
@@ -3209,10 +3220,18 @@ class llama_io_read_socket : public llama_io_read_i {
             staging_pos += take;
             done        += take;
             if (staging_pos >= staging_len && done < size) {
-                // recv the next chunk WHILE the H2D copy of the current one runs
-                ssize_t r = ::recv(fd, other.data(), other.size(), 0);
+                // recv the next chunk WHILE the H2D copy of the current one runs.
+                // hydra#713 review (finding 1): this pipelined branch is the
+                // dominant CUDA path (tensor_backend != nullptr) and carries the
+                // whole ~800 MB stream — it must retry EAGAIN/EWOULDBLOCK like
+                // refill() does, not treat the first -1 as stream-end.
+                ssize_t r = hydra_recv_with_retry(fd, other.data(), other.size(), 30000);
                 if (r <= 0) {
-                    throw std::runtime_error("hydra: socket recv failed during state restore");
+                    if (r == 0) {
+                        throw std::runtime_error("hydra: socket recv EOF during state restore");
+                    }
+                    throw std::runtime_error(std::string("hydra: socket recv failed during state restore: ")
+                                             + std::strerror(errno));
                 }
                 other_len = (size_t)r;
                 if (hash_state != nullptr) {

tests/CMakeLists.txt

@@ -182,6 +182,14 @@ target_link_libraries(test-hydra-rpc-stale-sock PRIVATE ggml)
 # #470: test uses ggml_backend_buffer_copy_tensor (declared in ggml-backend-impl.h).
 target_include_directories(test-hydra-rpc-stale-sock PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src)
 
+# hydra#713: EAGAIN/EWOULDBLOCK retry in hydra_recv_all and
+# llama_io_read_socket::refill(). Pure socket test — no model/GPU needed.
+# Verifies that recv returning EAGAIN is retried with poll (not treated as
+# EOF) and that true EOF / timeout still fail cleanly.
+if (NOT WIN32)
+    llama_build_and_test(test-hydra-recv-eagain.cpp)
+endif()
+
 if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries)
     llama_build_and_test(test-sampling.cpp)