Set up comment negotiation system

ChrisPenner · ChrisPenner · commit 3feb82765f8a · 2025-12-19T07:51:22.000-08:00
diff --git a/share-api/src/Share/Web/UCM/HistoryComments/Impl.hs b/share-api/src/Share/Web/UCM/HistoryComments/Impl.hs
@@ -1,8 +1,11 @@
 module Share.Web.UCM.HistoryComments.Impl (server) where
 
+import Control.Concurrent.STM.TBMQueue (TBMQueue, closeTBMQueue, newTBMQueueIO, readTBMQueue, writeTBMQueue)
 import Control.Monad.Except
 import Control.Monad.Trans.Maybe
-import Data.Either (partitionEithers)
+import Data.Set qualified as Set
+import Data.Set.NonEmpty qualified as NESet
+import Ki.Unlifted qualified as Ki
 import Network.WebSockets.Connection
 import Share.IDs
 import Share.IDs qualified as IDs
@@ -18,8 +21,10 @@ import Share.Web.Authorization qualified as AuthZ
 import Share.Web.Errors (Unimplemented (Unimplemented), reportError, respondError)
 import Share.Web.UCM.HistoryComments.Queries qualified as Q
 import Unison.Debug qualified as Debug
+import Unison.Hash32 (Hash32)
 import Unison.Server.HistoryComments.API qualified as HistoryComments
-import Unison.Server.HistoryComments.Types (HistoryCommentChunk (..), UploadCommentsResponse (..))
+import Unison.Server.HistoryComments.Types (HistoryCommentDownloaderChunk (..), HistoryCommentUploaderChunk (..), UploadCommentsResponse (..))
+import Unison.Server.HistoryComments.Types qualified as Sync
 import Unison.Server.Types
 import Unison.Util.Websockets
 import UnliftIO
@@ -40,68 +45,109 @@ downloadHistoryCommentsStreamImpl _mayUserId _conn = do
   respondError Unimplemented
 
 -- Re-run the given STM action at most n times, collecting the results into a list.
--- If the action returns Nothing, stop early and return what has been collected so far, along with a Bool indicating whether the action was exhausted.
+-- If the action returns Nothing, stop and return what has been collected so far, along with a Bool indicating whether the action was exhausted.
 fetchChunk :: (Show a) => Int -> STM (Maybe a) -> STM ([a], Bool)
 fetchChunk size action = do
   let go 0 = pure ([], False)
       go n = do
         optional action >>= \case
           Nothing -> do
             -- No more values available at the moment
-            pure ([], False)
+            empty
           Just Nothing -> do
             -- Queue is closed
             pure ([], True)
           Just (Just val) -> do
             Debug.debugM Debug.Temp "Fetched value from queue" val
-            (rest, exhausted) <- go (n - 1)
+            (rest, exhausted) <- go (n - 1) <|> pure ([], False)
             pure (val : rest, exhausted)
   go size
 
 uploadHistoryCommentsStreamImpl :: Maybe UserId -> BranchRef -> Connection -> WebApp ()
 uploadHistoryCommentsStreamImpl mayCallerUserId br@(BranchRef branchRef) conn = do
-  Debug.debugM Debug.Temp "uploadHistoryCommentsStreamImpl called with branchRef: " (IDs.toText branchRef, mayCallerUserId)
   callerUserId <- AuthN.requireAuthenticatedUser' mayCallerUserId
-  result <- withQueues @(MsgOrError Void UploadCommentsResponse) @(MsgOrError Void HistoryCommentChunk) wsMessageBufferSize wsMessageBufferSize conn \q@(Queues {receive}) -> runExceptT $ do
+  result <- withQueues @(MsgOrError UploadCommentsResponse HistoryCommentDownloaderChunk) @(MsgOrError Void HistoryCommentUploaderChunk) wsMessageBufferSize wsMessageBufferSize conn \q@(Queues {receive, send}) -> Ki.scoped \scope -> runExceptT $ do
     projectBranchSH@ProjectBranchShortHand {userHandle, projectSlug, contributorHandle} <- case IDs.fromText @ProjectBranchShortHand branchRef of
       Left err -> handleErrInQueue q (UploadCommentsGenericFailure $ IDs.toText err)
       Right pbsh -> pure pbsh
     let projectSH = ProjectShortHand {userHandle, projectSlug}
     mayInfo <- lift . runMaybeT $ mapMaybeT PG.runTransaction $ do
       project <- MaybeT $ PGQ.projectByShortHand projectSH
-      Debug.debugM Debug.Temp "FOUND PROJECT" (project)
       branch <- MaybeT $ PGQ.branchByProjectBranchShortHand projectBranchSH
-      Debug.debugM Debug.Temp "FOUND BRANCH" (branch)
       contributorUser <- for contributorHandle (MaybeT . UserQ.userByHandle)
-      Debug.debugM Debug.Temp "FOUND Contributor" (contributorUser)
       pure (project, branch, contributorUser)
     (project, _branch, contributorUser) <- maybe (handleErrInQueue q $ UploadCommentsProjectBranchNotFound br) pure $ mayInfo
-    authZ <-
+    !authZ <-
       lift (AuthZ.checkUploadToProjectBranchCodebase callerUserId project.projectId (user_id <$> contributorUser)) >>= \case
         Left _authErr -> handleErrInQueue q (UploadCommentsNotAuthorized br)
         Right authZ -> pure authZ
-    let loop :: ExceptT UploadCommentsResponse WebApp ()
-        loop = do
-          (chunk, closed) <- atomically $ fetchChunk insertCommentBatchSize do
-            receive <&> fmap \case
-              Msg (HistoryCommentErrorChunk err) -> (Left $ UploadCommentsGenericFailure err)
-              Msg chunk -> (Right chunk)
-              DeserialiseFailure msg -> (Left $ UploadCommentsGenericFailure msg)
-              UserErr err -> absurd err
-          Debug.debugM Debug.Temp "Processing chunk of size" (length chunk)
-          let (errs, chunks) = partitionEithers chunk
-          when (not $ null errs) $ Debug.debugM Debug.Temp "Got errors in chunk" (errs)
-          lift $ PG.runTransaction $ Q.insertHistoryComments authZ project.projectId chunks
-          for errs $ \err -> handleErrInQueue q err
-          when (not closed) loop
-    loop
+    hashesToCheckQ <- liftIO $ newTBMQueueIO 100
+    commentsQ <- liftIO $ newTBMQueueIO 100
+    errMVar <- liftIO newEmptyTMVarIO
+    _receiverThread <- lift $ Ki.fork scope $ receiverWorker receive errMVar hashesToCheckQ commentsQ
+    inserterThread <- lift $ Ki.fork scope $ inserterWorker authZ commentsQ project.projectId
+    _hashCheckingThread <- lift $ Ki.fork scope $ hashCheckingWorker send hashesToCheckQ
+    -- The inserter thread will finish when the client closes the connection.
+    atomically $ Ki.await inserterThread
   case result of
     Left err -> reportError err
     Right (Left err, _leftovers) -> reportError err
     Right (Right (), _leftovers) -> pure ()
   where
+    inserterWorker ::
+      AuthZ.AuthZReceipt ->
+      TBMQueue (Either Sync.HistoryComment Sync.HistoryCommentRevision) ->
+      ProjectId ->
+      WebApp ()
+    inserterWorker authZ commentsQ projectId = do
+      let loop = do
+            (chunk, closed) <- atomically (fetchChunk insertCommentBatchSize (readTBMQueue commentsQ))
+            PG.runTransaction $ Q.insertHistoryComments authZ projectId chunk
+            when (not closed) loop
+      loop
+
+    hashCheckingWorker ::
+      (MsgOrError err HistoryCommentDownloaderChunk -> STM Bool) ->
+      TBMQueue Hash32 ->
+      WebApp ()
+    hashCheckingWorker send hashesToCheckQ = do
+      let loop = do
+            (hashes, closed) <- atomically (fetchChunk insertCommentBatchSize (readTBMQueue hashesToCheckQ))
+            Debug.debugM Debug.Temp "Checking hashes chunk of size" (length hashes)
+            unknownHashes <- PG.runTransaction $ do Q.filterForUnknownHistoryCommentHashes hashes
+            case NESet.nonEmptySet (Set.fromList unknownHashes) of
+              Nothing -> pure ()
+              Just unknownHashesSet -> do
+                void . atomically $ send $ Msg $ RequestCommentsChunk unknownHashesSet
+            when (not closed) loop
+      loop
+    receiverWorker :: STM (Maybe (MsgOrError Void HistoryCommentUploaderChunk)) -> TMVar Text -> TBMQueue Hash32 -> TBMQueue (Either Sync.HistoryComment Sync.HistoryCommentRevision) -> WebApp ()
+    receiverWorker recv errMVar hashesToCheckQ commentsQ = do
+      let loop = do
+            next <- atomically do
+              recv >>= \case
+                Nothing -> do
+                  closeTBMQueue hashesToCheckQ
+                  closeTBMQueue commentsQ
+                  pure (pure ())
+                Just (DeserialiseFailure err) -> do
+                  putTMVar errMVar err
+                  pure (pure ())
+                Just (Msg msg) -> do
+                  case msg of
+                    PossiblyNewHashesChunk hashesToCheck -> do
+                      for_ hashesToCheck $ \h -> writeTBMQueue hashesToCheckQ h
+                    DoneSendingHashesChunk -> do
+                      closeTBMQueue hashesToCheckQ
+                    HistoryCommentChunk comment -> do
+                      writeTBMQueue commentsQ (Left comment)
+                    HistoryCommentRevisionChunk revision -> do
+                      writeTBMQueue commentsQ (Right revision)
+                  pure loop
+            next
+      loop
     insertCommentBatchSize = 100
-    handleErrInQueue :: forall o x e. Queues (MsgOrError e UploadCommentsResponse) o -> UploadCommentsResponse -> ExceptT UploadCommentsResponse WebApp x
+    handleErrInQueue :: forall o x e a. Queues (MsgOrError e a) o -> e -> ExceptT e WebApp x
     handleErrInQueue Queues {send} e = do
-      _ <- atomically $ send $ Msg e
+      _ <- atomically $ send $ UserErr e
       throwError e
diff --git a/share-api/src/Share/Web/UCM/HistoryComments/Queries.hs b/share-api/src/Share/Web/UCM/HistoryComments/Queries.hs
@@ -4,6 +4,7 @@
 module Share.Web.UCM.HistoryComments.Queries
   ( fetchProjectBranchCommentsSince,
     insertHistoryComments,
+    filterForUnknownHistoryCommentHashes,
   )
 where
 
@@ -25,7 +26,7 @@ import Share.Web.Authorization (AuthZReceipt)
 import Unison.Hash32 (Hash32)
 import Unison.Server.HistoryComments.Types
 
-fetchProjectBranchCommentsSince :: AuthZReceipt -> ProjectId -> CausalId -> UTCTime -> PG.Transaction e (PGCursor HistoryCommentChunk)
+fetchProjectBranchCommentsSince :: AuthZReceipt -> ProjectId -> CausalId -> UTCTime -> PG.Transaction e (PGCursor HistoryCommentUploaderChunk)
 fetchProjectBranchCommentsSince !_authZ projectId causalId sinceTime = do
   PG.newRowCursor @(Bool, Maybe Text, Maybe Text, Maybe Int64, Maybe Bool, Maybe ByteString, Maybe Hash32, Maybe Hash32, Maybe Text, Maybe Int64, Maybe Text, Maybe Hash32, Maybe Hash32)
     "fetchProjectBranchCommentsSince"
@@ -96,7 +97,7 @@ utcTimeToMillis utcTime =
     & (* (1_000 :: Rational))
     & round
 
-insertHistoryComments :: AuthZReceipt -> ProjectId -> [HistoryCommentChunk] -> PG.Transaction e ()
+insertHistoryComments :: AuthZReceipt -> ProjectId -> [Either HistoryComment HistoryCommentRevision] -> PG.Transaction e ()
 insertHistoryComments !_authZ projectId chunks = PG.pipelined $ do
   let thumbprints = NESet.nonEmptySet $ Set.fromList (comments <&> \HistoryComment {authorThumbprint} -> authorThumbprint)
   for thumbprints insertThumbprints
@@ -107,9 +108,8 @@ insertHistoryComments !_authZ projectId chunks = PG.pipelined $ do
   where
     (comments, revisions) =
       chunks & foldMap \case
-        HistoryCommentChunk comment -> ([comment], [])
-        HistoryCommentRevisionChunk revision -> ([], [revision])
-        HistoryCommentErrorChunk err -> error $ "HistoryCommentErrorChunk: " <> show err -- TODO Handle this
+        Left comment -> ([comment], [])
+        Right revision -> ([], [revision])
     insertHistoryComments :: [HistoryComment] -> PG.Pipeline e ()
     insertHistoryComments comments = do
       PG.execute_
@@ -201,3 +201,12 @@ insertHistoryComments !_authZ projectId chunks = PG.pipelined $ do
               commentHash,
               utcTimeToMillis createdAt
             )
+
+filterForUnknownHistoryCommentHashes :: (PG.QueryA m) => [Hash32] -> m [Hash32]
+filterForUnknownHistoryCommentHashes hashes = do
+  -- error "TODO: Check whether they're in the project as well."
+  PG.queryListCol
+    [PG.sql|
+      SELECT hash FROM ^{PG.singleColumnTable hashes} AS t(hash)
+        WHERE hash NOT IN (SELECT comment_hash FROM history_comments)
+    |]
