From 2d08ba778fd13ac20488440ffc3b3ddc7bbb4e53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Henry?= <gregoire.henry@ocamlpro.com>
Date: Sat, 11 Nov 2017 03:34:12 +0100
Subject: [PATCH] Utils: fix `Lwt_utils.with_timeout`

The function will never fail with `Timeout`, if the concurrent
computation terminates at the same "time".
---
 src/utils/lwt_utils.ml | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/utils/lwt_utils.ml b/src/utils/lwt_utils.ml
index 33358b6f0..891b2d75d 100644
--- a/src/utils/lwt_utils.ml
+++ b/src/utils/lwt_utils.ml
@@ -483,17 +483,17 @@ let () =
     (fun () -> Timeout)
 
 let with_timeout ?(canceler = Canceler.create ()) timeout f =
-  let t = Lwt_unix.sleep timeout in
-  Lwt.choose [
-    (t >|= fun () -> None) ;
-    (f canceler >|= fun x -> Some x)
-  ] >>= function
-  | Some x when Lwt.state t = Lwt.Sleep ->
-      Lwt.cancel t ;
-      Lwt.return x
-  | _ ->
-      Canceler.cancel canceler >>= fun () ->
-      fail Timeout
+  let timeout = Lwt_unix.sleep timeout in
+  let target = f canceler in
+  Lwt.choose [ timeout ; (target >|= fun _ -> ()) ] >>= fun () ->
+  Lwt_unix.yield () >>= fun () ->
+  if Lwt.state target <> Lwt.Sleep then begin
+    Lwt.cancel timeout ;
+    target
+  end else begin
+    Canceler.cancel canceler >>= fun () ->
+    fail Timeout
+  end
 
 let unless cond f =
   if cond then Lwt.return () else f ()