@@ -5,9 +5,10 @@ module Share.Web.UCM.SyncStream.Impl (server) where
5
5
6
6
import Conduit
7
7
import Control.Concurrent.STM qualified as STM
8
+ import Control.Concurrent.STM.TBMQueue qualified as STM
8
9
import Control.Monad.Except (ExceptT (ExceptT ))
9
10
import Control.Monad.Trans.Except (runExceptT )
10
- import Data.ByteString.Lazy qualified as BL
11
+ import Data.Conduit.Combinators qualified as Conduit
11
12
import Servant
12
13
import Servant.Conduit (ConduitToSourceIO (.. ))
13
14
import Servant.Types.SourceT qualified as SourceT
@@ -28,6 +29,8 @@ import Share.Web.Errors
28
29
import Share.Web.UCM.Sync.HashJWT qualified as HashJWT
29
30
import Share.Web.UCM.SyncStream.Queries qualified as SSQ
30
31
import U.Codebase.Sqlite.Orphans ()
32
+ import U.Codebase.Sqlite.TempEntity (TempEntity )
33
+ import Unison.Debug qualified as Debug
31
34
import Unison.Hash32 (Hash32 )
32
35
import Unison.Share.API.Hash (HashJWTClaims (.. ))
33
36
import Unison.SyncV2.API qualified as SyncV2
@@ -78,34 +81,44 @@ downloadEntitiesStreamImpl mayCallerUserId (SyncV2.DownloadEntitiesRequest {caus
78
81
authZToken <- lift AuthZ. checkDownloadFromProjectBranchCodebase `whenLeftM` \ _err -> throwError (SyncV2. DownloadEntitiesNoReadPermission branchRef)
79
82
let codebaseLoc = Codebase. codebaseLocationForProjectBranchCodebase projectOwnerUserId contributorId
80
83
pure $ Codebase. codebaseEnv authZToken codebaseLoc
81
- q <- liftIO $ STM. newTBQueueIO 10
84
+ q <- liftIO $ STM. newTBMQueueIO 10
82
85
streamResults <- lift $ UnliftIO. toIO do
86
+ Debug. debugLogM Debug. Temp " Starting source Stream"
83
87
Codebase. runCodebaseTransaction codebase $ do
84
88
(_bhId, causalId) <- CausalQ. expectCausalIdsOf id (hash32ToCausalHash causalHash)
85
89
cursor <- SSQ. allSerializedDependenciesOfCausalCursor causalId
86
90
Cursor. foldBatched cursor 1000 \ batch -> do
87
- PG. transactionUnsafeIO $ STM. atomically $ STM. writeTBQueue q batch
91
+ Debug. debugLogM Debug. Temp " Source stream batch"
92
+ PG. transactionUnsafeIO $ STM. atomically $ STM. writeTBMQueue q batch
93
+ PG. transactionUnsafeIO $ STM. atomically $ STM. closeTBMQueue q
88
94
pure $ conduitToSourceIO do
89
95
handle <- liftIO $ Async. async streamResults
90
96
stream q handle
97
+ Conduit. .| ( Conduit. iterM \ case
98
+ EntityChunk {hash} -> Debug. debugM Debug. Temp " Chunk " hash
99
+ ErrorChunk err -> Debug. debugM Debug. Temp " Error " err
100
+ )
91
101
where
92
- stream :: STM. TBQueue (NonEmpty (Hash32 , ByteString )) -> Async. Async ( ) -> ConduitT () DownloadEntitiesChunk IO ()
93
- stream q async = do
102
+ stream :: STM. TBMQueue (NonEmpty (SyncV2. CBORBytes TempEntity , Hash32 )) -> ( Async. Async a ) -> ConduitT () DownloadEntitiesChunk IO ()
103
+ stream q handle = do
94
104
let loop :: ConduitT () DownloadEntitiesChunk IO ()
95
105
loop = do
96
- next <- liftIO . STM. atomically $ do
97
- STM. tryReadTBQueue q >>= \ case
98
- Nothing -> do
99
- Async. waitSTM async $> Nothing
100
- Just batch -> do
101
- pure $ Just batch
102
- case next of
103
- Nothing -> pure ()
106
+ Debug. debugLogM Debug. Temp " Waiting for batch..."
107
+ liftIO (STM. atomically (STM. readTBMQueue q)) >>= \ case
108
+ -- The queue is closed.
109
+ Nothing -> do
110
+ Debug. debugLogM Debug. Temp " Queue closed. finishing up!"
111
+ pure ()
104
112
Just batch -> do
105
- let chunks = batch <&> \ (hash, bytes) -> EntityChunk {hash, entityCBOR = SyncV2. CBORBytes $ BL. fromStrict bytes}
113
+ let chunks = batch <&> \ (entityCBOR, hash) -> EntityChunk {hash, entityCBOR}
114
+ Debug. debugLogM Debug. Temp $ " Emitting chunk of " <> show (length chunks) <> " entities"
106
115
yieldMany chunks
107
116
loop
108
117
loop
118
+ Debug. debugLogM Debug. Temp " Waiting for worker thread to finish"
119
+ -- It _should_ have terminated by now, but just in case, cancel it.
120
+ Async. cancel handle
121
+ Debug. debugLogM Debug. Temp " Done!"
109
122
110
123
emitErr :: SyncV2. DownloadEntitiesError -> SourceIO SyncV2. DownloadEntitiesChunk
111
124
emitErr err = SourceT. source [ErrorChunk err]
0 commit comments