Discovery: track our search requests per peer

2 files changed, 119 insertions, 23 deletions

diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index 6e39b8c..d6bff83 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -26,6 +26,8 @@ import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Word
 
+import System.Clock
+
 import Text.Read
 
 #ifdef ENABLE_ICE_SUPPORT
@@ -209,7 +211,8 @@ emptyPeer = DiscoveryPeer
     }
 
 data DiscoveryPeerState = DiscoveryPeerState
-    { dpsOurTunnelRequests :: [ ( RefDigest, StreamWriter ) ]
+    { dpsWeAskedFor :: Map RefDigest SearchStatus
+    , dpsOurTunnelRequests :: [ ( RefDigest, StreamWriter ) ]
     -- ( original target, our write stream )
     , dpsRelayedTunnelRequests :: [ ( RefDigest, ( StreamReader, StreamWriter )) ]
     -- ( original source, ( from source, to target ))
@@ -223,6 +226,9 @@ data DiscoveryGlobalState = DiscoveryGlobalState
     , dgsSearchingFor :: Set RefDigest
     }
 
+data SearchStatus
+    = SearchingSince TimeSpec
+
 instance Service DiscoveryService where
     serviceID _ = mkServiceID "dd59c89c-69cc-4703-b75b-4ddcd4b3c23c"
 
@@ -231,7 +237,8 @@ instance Service DiscoveryService where
 
     type ServiceState DiscoveryService = DiscoveryPeerState
     emptyServiceState _ = DiscoveryPeerState
-        { dpsOurTunnelRequests = []
+        { dpsWeAskedFor = M.empty
+        , dpsOurTunnelRequests = []
         , dpsRelayedTunnelRequests = []
         , dpsStunServer = Nothing
         , dpsTurnServer = Nothing
@@ -311,18 +318,29 @@ instance Service DiscoveryService where
                 " for " <> show (either refDigest id edgst) <>
                 " result [" <> T.unpack (T.intercalate "," $ map toText results) <> "]"
 
-        DiscoveryResult _ [] -> do
-            -- not found
-            return ()
-
         DiscoveryResult edgst addrs -> do
             let dgst = either refDigest id edgst
-            -- TODO: check if we really requested that
             server <- asks svcServer
             st <- getStorage
             self <- svcSelf
             discoveryPeer <- asks svcPeer
-            let runAsService = runPeerService @DiscoveryService discoveryPeer
+            pid <- asks svcPeerIdentity
+
+            weAskedFor <- dpsWeAskedFor <$> svcGet
+            let askedFor = M.member dgst weAskedFor
+            debugLog $
+                "result from " <> show (refDigest $ storedRef $ idData pid) <>
+                " for " <> show dgst <> ": [" <> T.unpack (T.intercalate "," $ map toText addrs) <> "]" <>
+                (if askedFor then "" else " (not asked for)")
+
+            when askedFor $ do
+                let weAskedFor' = M.delete dgst weAskedFor
+                svcModify $ \s -> s { dpsWeAskedFor = weAskedFor' }
+                debugLog $
+                    "remains asked " <> show (refDigest $ storedRef $ idData pid) <>
+                    " for " <> show (M.keys weAskedFor')
+
+            let runAsService = runPeerService @DiscoveryService @IO discoveryPeer
 
             let tryAddresses = \case
                     DiscoveryIP ipaddr port : _ -> do
@@ -367,12 +385,13 @@ instance Service DiscoveryService where
                         discoverySetupTunnelResponse dgst
 
                     addr : rest -> do
-                        svcPrint $ "Discovery: unsupported address in result: " ++ T.unpack (toText addr)
+                        debugLog $ "unsupported address in result: " ++ T.unpack (toText addr)
                         tryAddresses rest
 
-                    [] -> svcPrint $ "Discovery: no (supported) address received for " <> show dgst
+                    [] -> debugLog $ "no (supported) address received for " <> show dgst
 
-            tryAddresses addrs
+            when askedFor $ do
+                tryAddresses addrs
 
         DiscoveryConnectionRequest conn -> do
             self <- svcSelf
@@ -550,8 +569,17 @@ instance Service DiscoveryService where
 
         when (not $ null addrs) $ do
             sendToPeer peer $ DiscoverySelf addrs Nothing
-        forM_ searchingFor $ \dgst -> do
-            sendToPeer peer $ DiscoverySearch (Right dgst)
+
+        when (not $ null searchingFor) $ do
+            forM_ searchingFor $ \dgst -> do
+                sendToPeer peer $ DiscoverySearch (Right dgst)
+
+            now <- liftIO $ getTime Monotonic
+            let weAskedFor' = M.fromAscList $ map (, SearchingSince now) $ S.toAscList searchingFor
+            svcModify $ \s -> s { dpsWeAskedFor = weAskedFor' }
+            debugLog $
+                "we asked new peer " <> show (refDigest $ storedRef $ idData pid) <>
+                " for " <> show (M.keys weAskedFor')
 
     serviceUpdatedPeer = do
         pid <- asks svcPeerIdentity
@@ -622,16 +650,20 @@ discoverySearch server dgst = do
                     return $ dgst `elem` identityDigests pid
                 _ -> return False
         when (not $ or match) $ do
-            alreadySearching <- modifyServiceGlobalState server (Proxy @DiscoveryService) $ \s ->
-                let alreadySearching = S.member dgst $ dgsSearchingFor s
-                 in ( if alreadySearching then s else s
-                         { dgsSearchingFor =  S.insert dgst $ dgsSearchingFor s
-                         }
-                    , alreadySearching
-                    )
-            when (not alreadySearching) $ do
-                forM_ peers $ \peer -> do
-                    sendToPeer peer $ DiscoverySearch $ Right dgst
+            _ <- modifyServiceGlobalState server (Proxy @DiscoveryService) $ \s ->
+                ( s { dgsSearchingFor =  S.insert dgst $ dgsSearchingFor s }, () )
+            now <- liftIO $ getTime Monotonic
+            forM_ peers $ \peer -> do
+                runPeerService peer $ do
+                    weAskedFor <- dpsWeAskedFor <$> svcGet
+                    when (not $ M.member dgst weAskedFor) $ do
+                        let weAskedFor' = M.insert dgst (SearchingSince now) weAskedFor
+                        svcModify $ \s -> s { dpsWeAskedFor = weAskedFor' }
+                        pid <- asks svcPeerIdentity
+                        debugLog $
+                            "we asked " <> show (refDigest $ storedRef $ idData pid) <>
+                            " for " <> show dgst <> " " <> show (M.keys weAskedFor')
+                        replyPacket $ DiscoverySearch $ Right dgst
 
 
 data TunnelAddress = TunnelAddress
diff --git a/test/discovery.et b/test/discovery.et
index 98043b9..ef5bbd4 100644
--- a/test/discovery.et
+++ b/test/discovery.et
@@ -92,6 +92,70 @@ test ManualDiscovery:
             expect /stop-server-done/ from p
 
 
+test DiscoveryJoinedAfter:
+    let services = "discovery"
+
+    subnet sd
+    subnet s1
+    subnet s2
+
+    spawn as pd on sd
+    spawn as p1 on s1
+    spawn as p2 on s2
+    send "create-identity Discovery" to pd
+    send "create-identity Device1 Owner1" to p1
+    send "create-identity Device2 Owner2" to p2
+
+    expect /create-identity-done ref ($refpat).*/ from p1 capture p1id
+    send "identity-info $p1id" to p1
+    expect /identity-info ref $p1id base ($refpat) owner ($refpat).*/ from p1 capture p1base, p1owner
+    send "identity-info $p1owner" to p1
+    expect /identity-info ref $p1owner base ($refpat).*/ from p1 capture p1obase
+
+    expect /create-identity-done ref $refpat.*/ from p2
+    expect /create-identity-done ref $refpat.*/ from pd
+
+    for id in [ p1obase ]:
+        for p in [ pd, p1, p2 ]:
+            send "start-server services $services" to p
+
+        for p in [ p2 ]:
+            with p:
+                send "peer-add ${pd.node.ip}"
+                expect:
+                    /peer 1 addr ${pd.node.ip} 29665/
+                    /peer 1 id Discovery/
+            expect from pd:
+                    /peer [12] addr ${p.node.ip} 29665/
+                    /peer [12] id Device. Owner./
+
+        send "discovery-connect $id" to p2
+
+        for p in [ p1 ]:
+            with p:
+                send "peer-add ${pd.node.ip}"
+                expect:
+                    /peer 1 addr ${pd.node.ip} 29665/
+                    /peer 1 id Discovery/
+            expect from pd:
+                    /peer [12] addr ${p.node.ip} 29665/
+                    /peer [12] id Device. Owner./
+
+        send "discovery-connect $id" to p2
+
+        expect from p1:
+            /peer [0-9]+ addr ${p2.node.ip} 29665/
+            /peer [0-9]+ id Device2 Owner2/
+        expect from p2:
+            /peer [0-9]+ addr ${p1.node.ip} 29665/
+            /peer [0-9]+ id Device1 Owner1/
+
+        for p in [ pd, p1, p2 ]:
+            send "stop-server" to p
+        for p in [ pd, p1, p2 ]:
+            expect /stop-server-done/ from p
+
+
 test DiscoveryTunnel:
     let services = "discovery:tunnel"