add ReadStream RPC (#1500)

* add readStream service

* fix the timing of readerStopReading calls

* some refactor && add readSingleShardStream RPC

* update proto

Author: YangKian

common/hstream/HStream/Exception.hs

@@ -50,6 +50,7 @@ module HStream.Exception
   , InvalidQuerySql (InvalidQuerySql)
   , SQLNotSupportedByParseSQL(SQLNotSupportedByParseSQL)
   , ConflictShardReaderOffset (ConflictShardReaderOffset)
+  , TooManyShardCount (TooManyShardCount)
 
     -- * Exception: SomeDeadlineExceeded
     --
@@ -427,6 +428,7 @@ MAKE_EX_1_DEFMSG(SomeInvalidArgument, InvalidConnectorType, Text, API.ErrorCodeC
 MAKE_EX_1_DEFMSG(SomeInvalidArgument, SQLNotSupportedByParseSQL, Text, API.ErrorCodeInternalError)
 MAKE_EX_1_DEFMSG(SomeInvalidArgument, InvalidQuerySql, String, API.ErrorCodeQueryInvalidSQL)
 MAKE_EX_1_DEFMSG(SomeInvalidArgument, ConflictShardReaderOffset, String, API.ErrorCodeShardReaderConflictOffset)
+MAKE_EX_1_DEFMSG(SomeInvalidArgument, TooManyShardCount, String, API.ErrorCodeShardReaderTooManyShards)
 
 invalidIdentifier :: API.ResourceType -> String -> InvalidObjectIdentifier
 invalidIdentifier API.ResourceTypeResStream       = InvalidObjectIdentifier API.ErrorCodeStreamInvalidObjectIdentifier

external/protocol

@@ -110,6 +110,7 @@ module HStream.Store.Stream
   , LD.readerSetWaitOnlyWhenNoData
   , LD.readerStopReading
   , LD.readerIsReadingAny
+  , LD.readerIsReading
     -- ** Checkpointed Reader
   , FFI.LDSyncCkpReader
   , newLDFileCkpReader

hstream-store/HStream/Store/Stream.hs

@@ -59,11 +59,13 @@ handlers serverContext@ServerContext{..} =
         -- Shards
         hstreamApiListShards = listShardsHandler serverContext,
         -- Reader
-        hstreamApiListShardReaders = listShardReadersHandler serverContext,
+        hstreamApiListShardReaders  = listShardReadersHandler serverContext,
         hstreamApiCreateShardReader = createShardReaderHandler serverContext,
         hstreamApiDeleteShardReader = deleteShardReaderHandler serverContext,
         hstreamApiReadShard         = readShardHandler serverContext,
         hstreamApiReadShardStream   = readShardStreamHandler serverContext,
+        hstreamApiReadStream        = readStreamHandler serverContext,
+        hstreamApiReadSingleShardStream = readSingleShardStreamHandler serverContext,
 
         -- Stats
         hstreamApiPerStreamTimeSeriesStats = H.perStreamTimeSeriesStats scStatsHolder,

hstream/src/HStream/Server/Core/ShardReader.hs

@@ -13,12 +13,16 @@ module HStream.Server.Handler.ShardReader
   , readShardHandler
   , listShardReadersHandler
   , readShardStreamHandler
+  , readStreamHandler
+  , readSingleShardStreamHandler
     -- * For hs-grpc-server
   , handleListShardReaders
   , handleCreateShardReader
   , handleDeleteShardReader
   , handleReadShard
   , handleReadShardStream
+  , handleReadStream
+  , handleReadSingleShardStream
   )
 where
 
@@ -122,3 +126,41 @@ handleReadShardStream
 handleReadShardStream sc _ req stream = catchDefaultEx $ do
   Log.debug $ "Receive read shard stream Request: " <> Log.build (show req)
   C.readShardStream sc req (G.streamWrite stream . Just)
+
+readStreamHandler
+  :: ServerContext
+  -> ServerRequest 'ServerStreaming ReadStreamRequest ReadStreamResponse
+  -> IO (ServerResponse 'ServerStreaming ReadStreamResponse)
+readStreamHandler sc (ServerWriterRequest _meta req streamSend) =
+  defaultServerStreamExceptionHandle $ do
+    Log.debug $ "Receive read shard stream Request: " <> Log.build (show req)
+    C.readStream sc req streamWrite
+    return $ ServerWriterResponse mempty StatusUnknown "should not reach here"
+  where
+    streamWrite x = first show <$> streamSend x
+
+handleReadStream
+  :: ServerContext
+  -> G.ServerStreamHandler ReadStreamRequest ReadStreamResponse ()
+handleReadStream sc _ req stream = catchDefaultEx $ do
+  Log.debug $ "Receive read shard stream Request: " <> Log.build (show req)
+  C.readStream sc req (G.streamWrite stream . Just)
+
+readSingleShardStreamHandler
+  :: ServerContext
+  -> ServerRequest 'ServerStreaming ReadSingleShardStreamRequest ReadSingleShardStreamResponse
+  -> IO (ServerResponse 'ServerStreaming ReadSingleShardStreamResponse)
+readSingleShardStreamHandler sc (ServerWriterRequest _meta req streamSend) =
+  defaultServerStreamExceptionHandle $ do
+    Log.debug $ "Receive read single shard stream Request: " <> Log.build (show req)
+    C.readSingleShardStream sc req streamWrite
+    return $ ServerWriterResponse mempty StatusUnknown "should not reach here"
+  where
+    streamWrite x = first show <$> streamSend x
+
+handleReadSingleShardStream
+  :: ServerContext
+  -> G.ServerStreamHandler ReadSingleShardStreamRequest ReadSingleShardStreamResponse ()
+handleReadSingleShardStream sc _ req stream = catchDefaultEx $ do
+  Log.debug $ "Receive read single shard stream Request: " <> Log.build (show req)
+  C.readSingleShardStream sc req (G.streamWrite stream . Just)

hstream/src/HStream/Server/Handler.hs

@@ -55,6 +55,8 @@ handlers sc =
     -- Read
   , unary (GRPC :: GRPC P.HStreamApi "readShard") (H.handleReadShard sc)
   , serverStream (GRPC :: GRPC P.HStreamApi "readShardStream") (H.handleReadShardStream sc)
+  , serverStream (GRPC :: GRPC P.HStreamApi "readStream") (H.handleReadStream sc)
+  , serverStream (GRPC :: GRPC P.HStreamApi "readSingleShardStream") (H.handleReadSingleShardStream sc)
     -- Subscribe
   , bidiStream (GRPC :: GRPC P.HStreamApi "streamingFetch") (H.handleStreamingFetch sc)
     -- Stats

hstream/src/HStream/Server/Handler/ShardReader.hs

@@ -12,6 +12,7 @@ import           Control.Concurrent               (MVar, ThreadId)
 import           Control.Concurrent.STM
 import           Data.Aeson                       (FromJSON (..), ToJSON (..))
 import qualified Data.Aeson                       as Aeson
+import           Data.HashMap.Strict              (HashMap)
 import qualified Data.HashMap.Strict              as HM
 import qualified Data.Heap                        as Heap
 import           Data.Int                         (Int32, Int64)
@@ -255,14 +256,25 @@ instance TaskManager IO.Worker where
   recoverTask = IO.recoverTask
 
 data ShardReader = ShardReader
-  { reader       :: S.LDReader
-  , totalBatches :: Maybe (IORef Word64)
-  , startTs      :: Maybe Int64
-  , endTs        :: Maybe Int64
+  { shardReader             :: S.LDReader
+  , shardReaderTotalBatches :: Maybe (IORef Word64)
+  , shardReaderStartTs      :: Maybe Int64
+  , shardReaderEndTs        :: Maybe Int64
+  , targetShard             :: S.C_LogID
   }
 
-mkShardReader :: S.LDReader ->  Maybe (IORef Word64) -> Maybe Int64 -> Maybe Int64 -> ShardReader
-mkShardReader reader totalBatches startTs endTs = ShardReader {..}
+mkShardReader :: S.LDReader -> S.C_LogID -> Maybe (IORef Word64) -> Maybe Int64 -> Maybe Int64 -> ShardReader
+mkShardReader shardReader targetShard shardReaderTotalBatches shardReaderStartTs shardReaderEndTs = ShardReader {..}
+
+data StreamReader = StreamReader
+  { streamReader             :: S.LDReader
+  , streamReaderTotalBatches :: Maybe (IORef Word64)
+  , streamReaderTsLimits     :: HashMap S.C_LogID (Maybe Int64, Maybe Int64)
+    -- ^ shardId -> (startTs, endTs)
+  }
+
+mkStreamReader :: S.LDReader ->  Maybe (IORef Word64) -> HashMap S.C_LogID (Maybe Int64, Maybe Int64) -> StreamReader
+mkStreamReader streamReader streamReaderTotalBatches streamReaderTsLimits = StreamReader {..}
 
 --------------------------------------------------------------------------------