implement split/merge shard

YangKian · YangKian · commit 002bc55c4d4e · 2022-07-20T10:54:59.000+08:00
diff --git a/hstream/src/HStream/Server/Core/Shard.hs b/hstream/src/HStream/Server/Core/Shard.hs
@@ -1,31 +1,42 @@
+{-# LANGUAGE LambdaCase #-}
+
 module HStream.Server.Core.Shard
-    ( readShard,
-      listShards,
-      splitShards,
-    )
+  ( readShard
+  , listShards
+  , splitShards
+  , mergeShards
+  )
 where
 
-import           Control.Exception                (bracket)
-import           Control.Monad                    (void)
-import           Data.Foldable                    (foldl')
-import qualified Data.Map.Strict                  as M
-import qualified Data.Vector                      as V
-import           GHC.Stack                        (HasCallStack)
-import           Network.GRPC.HighLevel.Generated
-import qualified Z.Data.CBytes                    as CB
-
-import           Data.Maybe                       (fromJust)
-import           HStream.Connector.HStore         (transToStreamName)
-import           HStream.Server.Exception         (InvalidArgument (..),
-                                                   StreamNotExist (..))
-import           HStream.Server.Handler.Common    (decodeRecordBatch)
-import qualified HStream.Server.HStreamApi        as API
-import           HStream.Server.ReaderPool        (getReader, putReader)
-import           HStream.Server.Types             (ServerContext (..))
-import qualified HStream.Store                    as S
-import           HStream.ThirdParty.Protobuf      as PB
+import           Control.Exception          (bracket)
+import           Control.Monad              (foldM, void)
+import           Data.Foldable              (foldl')
+import qualified Data.HashMap.Strict        as HM
+import qualified Data.Map.Strict            as M
+import qualified Data.Vector                as V
+import           GHC.Stack                  (HasCallStack)
+import qualified Z.Data.CBytes              as CB
+
+import           Control.Concurrent         (MVar, modifyMVar, modifyMVar_,
+                                             readMVar)
+import           Data.Maybe                 (fromJust, fromMaybe)
+import           Data.Text                  (Text)
+import           Data.Word                  (Word64)
+import           HStream.Server.Core.Common (decodeRecordBatch)
+import qualified HStream.Server.HStreamApi  as API
+import           HStream.Server.ReaderPool  (getReader, putReader)
+import           HStream.Server.Shard       (Shard (..), ShardKey (..),
+                                             SharedShardMap, cBytesToKey,
+                                             mergeTwoShard, mkShard,
+                                             mkSharedShardMapWithShards,
+                                             shardKeyToText, splitByKey,
+                                             splitHalf, textToShardKey)
+import           HStream.Server.Types       (ServerContext (..),
+                                             transToStreamName)
+import qualified HStream.Store              as S
 import           HStream.Utils
-import           Proto3.Suite                     (Enumerated (Enumerated))
+import           Proto3.Suite               (Enumerated (Enumerated))
+import           Z.Data.CBytes              (CBytes)
 
 -------------------------------------------------------------------------------
 
@@ -36,16 +47,13 @@ listShards
   -> IO (V.Vector API.Shard)
 listShards ServerContext{..} API.ListShardsRequest{..} = do
   shards <- M.elems <$> S.listStreamPartitions scLDClient streamId
-  V.foldM' getShardInfo V.empty $ V.fromList shards
+  V.foldM' constructShard V.empty $ V.fromList shards
  where
    streamId = transToStreamName listShardsRequestStreamName
-   startKey = CB.pack "startKey"
-   endKey   = CB.pack "endKey"
-   epoch    = CB.pack "epoch"
 
-   getShardInfo shards logId = do
+   constructShard shards logId = do
      attr <- S.getStreamPartitionExtraAttrs scLDClient logId
-     case getInfo attr of
+     case getShardInfo attr of
        Nothing -> return . V.snoc shards $
          API.Shard { API.shardStreamName = listShardsRequestStreamName
                    , API.shardShardId    = logId
@@ -54,19 +62,13 @@ listShards ServerContext{..} API.ListShardsRequest{..} = do
        Just(sKey, eKey, ep) -> return . V.snoc shards $
          API.Shard { API.shardStreamName        = listShardsRequestStreamName
                    , API.shardShardId           = logId
-                   , API.shardStartHashRangeKey = sKey
-                   , API.shardEndHashRangeKey   = eKey
+                   , API.shardStartHashRangeKey = shardKeyToText sKey
+                   , API.shardEndHashRangeKey   = shardKeyToText eKey
                    , API.shardEpoch             = ep
                    -- FIXME: neet a way to find if this shard is active
                    , API.shardIsActive          = True
                    }
 
-   getInfo mp = do
-     startHashRangeKey <- cBytesToText <$> M.lookup startKey mp
-     endHashRangeKey   <- cBytesToText <$> M.lookup endKey mp
-     shardEpoch        <- read . CB.unpack <$> M.lookup epoch mp
-     return (startHashRangeKey, endHashRangeKey, shardEpoch)
-
 readShard
   :: HasCallStack
   => ServerContext
@@ -101,8 +103,90 @@ readShard ServerContext{..} API.ReadShardRequest{..} = do
       return $ foldl' (\acc (_, _, _, record) -> acc <> record) V.empty receivedRecordsVecs
 
 splitShards
-  :: HasCallStack
-  => ServerContext
+  :: ServerContext
   -> API.SplitShardsRequest
   -> IO (V.Vector API.Shard)
-splitShards = error ""
+splitShards ServerContext{..} API.SplitShardsRequest{..} = do
+  sharedShardMp <- getShardMap scLDClient shardInfo splitShardsRequestStreamName
+  newShards <- splitShard sharedShardMp
+  updateShardTable newShards
+  return . V.map (shardToPb splitShardsRequestStreamName) $ V.fromList newShards
+ where
+   splitKey = textToShardKey splitShardsRequestSplitKey
+
+   split :: Bool -> ShardKey -> SharedShardMap -> IO (Shard, Shard)
+   split True  key mps = splitHalf scLDClient mps key
+   split False key mps = splitByKey scLDClient mps key
+
+   splitShard sharedShardMp =
+     modifyMVar shardInfo $ \info -> do
+       (s1, s2) <- split splitShardsRequestHalfSplit splitKey sharedShardMp
+       return (HM.insert splitShardsRequestStreamName sharedShardMp info, [s1, s2])
+
+   updateShardTable newShards =
+     modifyMVar_ shardTable $ \mp -> do
+       let dict = fromMaybe M.empty $ HM.lookup splitShardsRequestStreamName mp
+           dict' = foldl' (\acc Shard{startKey=sKey, shardId=sId} -> M.insert sKey sId acc) dict newShards
+       return $ HM.insert splitShardsRequestStreamName dict' mp
+
+mergeShards
+  :: ServerContext
+  -> API.MergeShardsRequest
+  -> IO API.Shard
+mergeShards ServerContext{..} API.MergeShardsRequest{..} = do
+  sharedShardMp <- getShardMap scLDClient shardInfo mergeShardsRequestStreamName
+  (newShard, removedKey) <- mergeShard sharedShardMp
+  updateShardTable newShard removedKey
+  return . shardToPb mergeShardsRequestStreamName $ newShard
+ where
+   mergeShard sharedShardMp = do
+     modifyMVar shardInfo $ \info -> do
+       let [shardKey1, shardKey2] = V.toList . V.map textToShardKey $ mergeShardsRequestShardKeys
+       res <- mergeTwoShard scLDClient sharedShardMp shardKey1 shardKey2
+       return (HM.insert mergeShardsRequestStreamName sharedShardMp info, res)
+
+   updateShardTable Shard{startKey=sKey, shardId=sId} removedKey =
+     modifyMVar_ shardTable $ \mp -> do
+       let dict = fromMaybe M.empty $ HM.lookup mergeShardsRequestStreamName mp
+           dict' = M.insert sKey sId dict
+           dict'' = M.delete removedKey dict'
+       return $ HM.insert mergeShardsRequestStreamName dict'' mp
+
+getShardMap :: S.LDClient -> MVar (HM.HashMap Text SharedShardMap) -> Text -> IO SharedShardMap
+getShardMap client shardInfo streamName = do
+  let streamId = transToStreamName streamName
+  readMVar shardInfo >>= pure <$> HM.lookup streamName >>= \case
+    Just mps -> return mps
+    Nothing  -> loadSharedShardMap client streamId
+
+loadSharedShardMap :: S.LDClient -> S.StreamId -> IO SharedShardMap
+loadSharedShardMap client streamId = do
+  shardIds <- M.elems <$> S.listStreamPartitions client streamId
+  mkSharedShardMapWithShards =<< foldM createShard [] shardIds
+ where
+   createShard acc shardId = do
+     attrs <- S.getStreamPartitionExtraAttrs client shardId
+     case getShardInfo attrs of
+       Nothing -> return acc
+       Just (sKey, eKey, epoch) -> return $ mkShard shardId streamId sKey eKey epoch : acc
+
+getShardInfo :: M.Map CBytes CBytes -> Maybe (ShardKey, ShardKey, Word64)
+getShardInfo mp = do
+  startHashRangeKey <- cBytesToKey <$> M.lookup startKey mp
+  endHashRangeKey   <- cBytesToKey <$> M.lookup endKey mp
+  shardEpoch        <- read . CB.unpack <$> M.lookup epoch mp
+  return (startHashRangeKey, endHashRangeKey, shardEpoch)
+ where
+   startKey = CB.pack "startKey"
+   endKey   = CB.pack "endKey"
+   epoch    = CB.pack "epoch"
+
+shardToPb :: Text -> Shard -> API.Shard
+shardToPb sName Shard{..} = API.Shard
+  { API.shardShardId           = shardId
+  , API.shardStreamName        = sName
+  , API.shardStartHashRangeKey = shardKeyToText startKey
+  , API.shardEndHashRangeKey   = shardKeyToText endKey
+  , API.shardEpoch             = epoch
+  , API.shardIsActive          = True
+  }
diff --git a/hstream/src/HStream/Server/Core/Stream.hs b/hstream/src/HStream/Server/Core/Stream.hs
@@ -27,16 +27,11 @@ import           Data.Text                         (Text)
 import qualified Data.Text                         as Text
 import qualified Data.Vector                       as V
 import           GHC.Stack                         (HasCallStack)
-import           Network.GRPC.HighLevel.Generated
-import           Proto3.Suite                      (Enumerated (Enumerated))
-import qualified Z.Data.CBytes                     as CB
 
 import           Control.Concurrent                (MVar, modifyMVar,
                                                     modifyMVar_)
-import           HStream.Connector.HStore          (transToStreamName)
 import           HStream.Common.ConsistentHashing  (getAllocatedNodeId)
 import qualified HStream.Logger                    as Log
-import           HStream.Server.Core.Common        (decodeRecordBatch)
 import           HStream.Server.Exception          (InvalidArgument (..),
                                                     StreamNotExist (..),
                                                     WrongServer (..))
@@ -54,13 +49,12 @@ import           HStream.Server.Types              (ServerContext (..),
                                                     transToStreamName)
 import qualified HStream.Stats                     as Stats
 import qualified HStream.Store                     as S
-import           HStream.ThirdParty.Protobuf       as PB
 import           HStream.Utils
 
 -------------------------------------------------------------------------------
 
 createStream :: HasCallStack => ServerContext -> API.Stream -> IO ()
-createStream ServerContext{..} stream@API.Stream{
+createStream ServerContext{..} API.Stream{
   streamBacklogDuration = backlogSec, streamShardCount = shardCount, ..} = do
 
   when (streamReplicationFactor == 0) $ throwIO (InvalidArgument "Stream replicationFactor cannot be zero")
diff --git a/hstream/src/HStream/Server/Exception.hs b/hstream/src/HStream/Server/Exception.hs
@@ -21,6 +21,8 @@ import           ZooKeeper.Exception
 
 import qualified HStream.Logger                       as Log
 import           HStream.Server.Persistence.Exception (PersistenceException)
+import           HStream.Server.Shard                 (CanNotMerge, CanNotSplit,
+                                                       ShardNotExist)
 import qualified HStream.Store                        as Store
 import           HStream.Utils                        (mkServerErrResp)
 
@@ -106,7 +108,7 @@ finalExceptionHandlers = [
     Log.fatal $ Log.buildString' err
     return (StatusInternal, mkStatusDetails err)
   ,
-  Handler $ \(err :: AsyncCancelled) -> do
+  Handler $ \(_ :: AsyncCancelled) -> do
     return (StatusOk, "")
   ,
   Handler $ \(err :: SomeException) -> do
@@ -144,6 +146,19 @@ zooKeeperExceptionHandler = [
   Handler $ \(e :: ZooException       ) -> handleZKException e StatusInternal
   ]
 
+shardExceptionHandler :: Handlers (StatusCode, StatusDetails)
+shardExceptionHandler = [
+  Handler $ \(e :: CanNotSplit) -> do
+    Log.fatal $ Log.buildString' e
+    return (StatusFailedPrecondition, mkStatusDetails e),
+  Handler $ \(e :: CanNotMerge) -> do
+    Log.fatal $ Log.buildString' e
+    return (StatusFailedPrecondition, mkStatusDetails e),
+  Handler $ \(e :: ShardNotExist) -> do
+    Log.fatal $ Log.buildString' e
+    return (StatusNotFound, mkStatusDetails e)
+  ]
+
 defaultHandlers :: Handlers (StatusCode, StatusDetails)
 defaultHandlers = serverExceptionHandlers
                ++ storeExceptionHandlers
diff --git a/hstream/src/HStream/Server/Handler.hs b/hstream/src/HStream/Server/Handler.hs
@@ -55,6 +55,7 @@ handlers serverContext@ServerContext{..} =
         hstreamApiReadShard = readShardHandler serverContext,
         hstreamApiListShards = listShardsHandler serverContext,
         hstreamApiSplitShards = splitShardsHandler serverContext,
+        hstreamApiMergeShards = mergeShardsHandler serverContext,
 
         -- Stats
         hstreamApiPerStreamTimeSeriesStats = H.perStreamTimeSeriesStats scStatsHolder,
diff --git a/hstream/src/HStream/Server/Handler/Shard.hs b/hstream/src/HStream/Server/Handler/Shard.hs
@@ -9,12 +9,15 @@ module HStream.Server.Handler.Shard
   ( listShardsHandler
   , readShardHandler
   , splitShardsHandler
+  , mergeShardsHandler
   )
 where
 
 import           Control.Exception
 import           Network.GRPC.HighLevel.Generated
 
+import           Control.Monad                    (when)
+import qualified Data.Vector                      as V
 import qualified HStream.Logger                   as Log
 import qualified HStream.Server.Core.Shard        as C
 import           HStream.Server.Exception
@@ -29,7 +32,7 @@ listShardsHandler
   :: ServerContext
   -> ServerRequest 'Normal ListShardsRequest ListShardsResponse
   -> IO (ServerResponse 'Normal ListShardsResponse)
-listShardsHandler sc (ServerNormalRequest _metadata request) = do
+listShardsHandler sc (ServerNormalRequest _metadata request) = defaultExceptionHandle $ do
   Log.debug "Receive List Shards Request"
   C.listShards sc request >>= returnResp . ListShardsResponse
 
@@ -45,14 +48,35 @@ splitShardsHandler
   :: ServerContext
   -> ServerRequest 'Normal SplitShardsRequest SplitShardsResponse
   -> IO (ServerResponse 'Normal SplitShardsResponse)
-splitShardsHandler sc (ServerNormalRequest _metadata request) = do
+splitShardsHandler sc (ServerNormalRequest _metadata request) = shardExceptionHandle $ do
   Log.debug $ "Receive Split Shards Request: " <> Log.buildString' (show request)
   C.splitShards sc request >>= returnResp . SplitShardsResponse
 
+mergeShardsHandler
+  :: ServerContext
+  -> ServerRequest 'Normal MergeShardsRequest MergeShardsResponse
+  -> IO (ServerResponse 'Normal MergeShardsResponse)
+mergeShardsHandler sc (ServerNormalRequest _metadata request@MergeShardsRequest{..}) = shardExceptionHandle $ do
+  Log.debug $ "Receive Merge Shards Request: " <> Log.buildString' (show request)
+  when (V.length mergeShardsRequestShardKeys /= 2) $ throwIO WrongShardCnt
+  C.mergeShards sc request >>= returnResp . MergeShardsResponse . Just
+
 -----------------------------------------------------------------------------------
 
 readShardExceptionHandle :: ExceptionHandle (ServerResponse 'Normal a)
 readShardExceptionHandle = mkExceptionHandle . setRespType mkServerErrResp $
   [ Handler (\(err :: Store.NOTFOUND) ->
       return (StatusUnavailable, mkStatusDetails err))
   ] ++ defaultHandlers
+
+shardExceptionHandle :: ExceptionHandle (ServerResponse 'Normal a)
+shardExceptionHandle = mkExceptionHandle . setRespType mkServerErrResp $
+  [ Handler (\(err :: WrongShardCnt) ->
+      return (StatusInvalidArgument, mkStatusDetails err))
+  ] ++ shardExceptionHandler ++ defaultHandlers
+
+data WrongShardCnt = WrongShardCnt
+instance Show WrongShardCnt where
+  show _ = "Only two shards can be merged at a time"
+
+instance Exception WrongShardCnt where
diff --git a/hstream/src/HStream/Server/Shard.hs b/hstream/src/HStream/Server/Shard.hs
