From 5f34796dc4a615c8fe496bbb9cc18a489bc5d107 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ludovic=20Court=C3=A8s?= <ludo@gnu.org>
Date: Tue, 23 Jan 2024 14:27:30 +0100
Subject: [PATCH] =?UTF-8?q?marionette:=20Add=20#:peek=3F=20to=20=E2=80=98w?=
 =?UTF-8?q?ait-for-tcp-port=3F=E2=80=99.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* gnu/build/marionette.scm (wait-for-tcp-port): Add #:peek? parameter
and honor it.

Change-Id: Ie7515a5223299390ab8af6fe5aa3cf63ba5c8078
---
 gnu/build/marionette.scm | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/gnu/build/marionette.scm b/gnu/build/marionette.scm
index 27c10e3dfe..0b0a8a70d8 100644
--- a/gnu/build/marionette.scm
+++ b/gnu/build/marionette.scm
@@ -1,5 +1,5 @@
 ;;; GNU Guix --- Functional package management for GNU
-;;; Copyright © 2016-2022 Ludovic Courtès <ludo@gnu.org>
+;;; Copyright © 2016-2022, 2024 Ludovic Courtès <ludo@gnu.org>
 ;;; Copyright © 2018 Chris Marusich <cmmarusich@gmail.com>
 ;;; Copyright © 2022, 2023 Maxim Cournoyer <maxim.cournoyer@gmail.com>
 ;;; Copyright © 2023 Bruno Victal <mirai@makinata.eu>
@@ -223,29 +223,49 @@ FILE has not shown up after TIMEOUT seconds, raise an error."
 (define* (wait-for-tcp-port port marionette
                             #:key
                             (timeout 20)
+                            (peek? #f)
                             (address `(make-socket-address AF_INET
                                                            INADDR_LOOPBACK
                                                            ,port)))
   "Wait for up to TIMEOUT seconds for PORT to accept connections in
 MARIONETTE.  ADDRESS must be an expression that returns a socket address,
-typically a call to 'make-socket-address'.  Raise an error on failure."
+typically a call to 'make-socket-address'.  When PEEK? is true, attempt to
+read a byte from the socket upon connection; retry if that gives the
+end-of-file object.
+
+Raise an error on failure."
   ;; Note: The 'connect' loop has to run within the guest because, when we
   ;; forward ports to the host, connecting to the host never raises
   ;; ECONNREFUSED.
   (match (marionette-eval
-          `(let* ((address ,address)
-                  (sock (socket (sockaddr:fam address) SOCK_STREAM 0)))
-             (let loop ((i 0))
+          `(let* ((address ,address))
+             (define (open-socket)
+               (socket (sockaddr:fam address) SOCK_STREAM 0))
+
+             (let loop ((sock (open-socket))
+                        (i 0))
                (catch 'system-error
                  (lambda ()
                    (connect sock address)
+                   (when ,peek?
+                     (let ((byte ((@ (ice-9 binary-ports) lookahead-u8)
+                                  sock)))
+                       (when (eof-object? byte)
+                         (close-port sock)
+                         (throw 'system-error
+                                "wait-for-tcp-port" "~A"
+                                (list (strerror ECONNRESET))
+                                (list ECONNRESET)))))
                    (close-port sock)
                    'success)
                  (lambda args
                    (if (< i ,timeout)
                        (begin
                          (sleep 1)
-                         (loop (+ 1 i)))
+                         (loop (if (port-closed? sock)
+                                   (open-socket)
+                                   sock)
+                               (+ 1 i)))
                        (list 'failure address))))))
           marionette)
     ('success #t)