From d5eaf0f567334b783912906d7dfe59640432c672 Mon Sep 17 00:00:00 2001
From: Joey Hess <joeyh@joeyh.name>
Date: Mon, 22 Jul 2024 16:56:44 -0400
Subject: [PATCH] improve clientKeepLocked

---
 Command/P2PHttp.hs | 11 ++++-------
 P2P/Http.hs        | 33 ++++++++++++++++++---------------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/Command/P2PHttp.hs b/Command/P2PHttp.hs
index 622dceffd1..e33e839f73 100644
--- a/Command/P2PHttp.hs
+++ b/Command/P2PHttp.hs
@@ -133,18 +133,15 @@ getAuthEnv = do
 testKeepLocked = do
 	mgr <- httpManager <$> getUrlOptions
 	burl <- liftIO $ parseBaseUrl "http://localhost:8080/"
-	keeplocked <- liftIO newEmptyTMVarIO
-	_ <- liftIO $ forkIO $ do
-		print "running, press enter to drop lock"
-		_ <- getLine
-		atomically $ writeTMVar keeplocked False
 	liftIO $ clientKeepLocked (mkClientEnv mgr burl)
 		(P2P.ProtocolVersion 3)
 		(B64UUID (toUUID ("lck" :: String)))
 		(B64UUID (toUUID ("cu" :: String)))
 		(B64UUID (toUUID ("su" :: String)))
-		[]
-		keeplocked
+		[] $ \keeplocked -> do
+			print "running, press enter to drop lock"
+			_ <- getLine
+			atomically $ writeTMVar keeplocked False
 
 testCheckPresent = do
 	mgr <- httpManager <$> getUrlOptions
diff --git a/P2P/Http.hs b/P2P/Http.hs
index 907fb59a78..205478c5a5 100644
--- a/P2P/Http.hs
+++ b/P2P/Http.hs
@@ -873,28 +873,31 @@ clientKeepLocked
 	-> B64UUID ClientSide
 	-> B64UUID ServerSide
 	-> [B64UUID Bypass]
-	-> TMVar Bool
+	-> (TMVar Bool -> IO ())
+	-- ^ The TMVar can be filled any number of times with True to send
+	-- repeated keep locked requests, eg to keep a connection alive.
+	-- Once filled with False, the lock will be dropped.
 	-> IO ()
-clientKeepLocked clientenv (ProtocolVersion ver) lckid cu su bypass keeplocked = do
+clientKeepLocked clientenv (ProtocolVersion ver) lckid cu su bypass a = do
+	keeplocked <- newEmptyTMVarIO
+	tid <- async $ a keeplocked
+	let cli' = cli lckid cu bypass
+		(Just connectionKeepAlive) (Just keepAlive)
+		(S.fromStepT (unlocksender keeplocked))
 	withClientM cli' clientenv $ \case
-		Left err  -> throwM err
 		Right (LockResult _ _) ->
-			liftIO $ print "end of lock connection to server"
+			wait tid
+		Left err  -> do
+			wait tid
+			throwM err
   where
-	unlocksender =
+	unlocksender keeplocked =
 		S.Yield (UnlockRequest False) $ S.Effect $ do
-			liftIO $ print "sent keep locked request"
 			return $ S.Effect $ do
 				stilllocked <- liftIO $ atomically $ takeTMVar keeplocked
-				if stilllocked
-					then return unlocksender
-					else do
-						liftIO $ print "sending unlock request"
-						return $ S.Yield (UnlockRequest True) S.Stop
-
-	cli' = cli lckid cu bypass
-		(Just connectionKeepAlive) (Just keepAlive)
-		(S.fromStepT unlocksender)
+				return $ if stilllocked
+					then unlocksender keeplocked
+					else S.Yield (UnlockRequest True) S.Stop
 	
 	cli = case ver of
 		3 -> v3 su V3