{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module Glean.Util.ShardManager
  ( ShardManager(..)
  , SomeShardManager(..)
  , BaseOfStack(..)
  , noSharding
  , shardByRepo
  , shardByBaseOfStackRepoHash
  , shardByRepoHash
  ) where

import Data.Int (Int64)
import Data.Map (Map)
import Data.Text (Text)
import qualified Glean.Types as Glean

-- | The base db of a stack (or a single db)
newtype BaseOfStack = BaseOfStack Glean.Repo

-- | An abstraction for sharding Glean DBs across multiple servers
data ShardManager shard = ShardManager
  {
    -- | Returns the list of shards for the current server.
    --   Nothing means fall back to no sharding.
    getAssignedShards :: IO (Maybe [shard])
    -- | Computes a mapping from DB stacks to shard identifiers.
    --   Takes as input the base of the stack.
    --   For non stacked dbs, this is just the db itself.
    --   See note [DB to Shard]
  , computeShardMapping :: IO (BaseOfStack -> Glean.Repo -> shard)
  , countersForShardSizes :: Map shard Int64 -> [(Text, Int)]
  }

data NoSharding = NoSharding
  deriving (Eq, Ord, Show)

-- | A sharding strategy with a single shard and trivial shard assignment
noSharding :: ShardManager NoSharding
noSharding = ShardManager
  (pure (Just [NoSharding]))
  (pure $ const $ const NoSharding)
  (const mempty)

-- | An existential wrapper around a 'ShardManager'
data SomeShardManager where
  SomeShardManager
    :: (Ord shard, Show shard) => ShardManager shard -> SomeShardManager

{- Note: [DB to Shard]

  'dbToShard' only wants the base of the stack. Why?

   We want the following invariant:

     > Every db in a stack maps to the same shard.

  This is for efficiency reasons, as the Glean server needs to download all
  the dependencies of a DB in order to be able to serve it.

  A simple way to enforce this invariant is to use the base of the stack for
  computing the shard id for all the dbs in the stack.
-}

-- | Shard by repo name using the provided dynamic shard assignment.
--   If the shard assignment is 'Nothing', fall back to no sharding
shardByRepo :: IO (Maybe [Text]) -> ShardManager Text
shardByRepo getAssignedShards =
  ShardManager {
    getAssignedShards = getAssignedShards,
    computeShardMapping = pure $ \(BaseOfStack db) _ -> Glean.repo_name db,
    countersForShardSizes = const []
  }

-- | A shard manager that uses repo hashes as shards,
-- and a dynamic shard assignment
shardByRepoHash :: IO (Maybe [Text]) -> ShardManager Text
shardByRepoHash refShardAssignment = ShardManager
  refShardAssignment
  (pure (\_ Glean.Repo{..} -> repo_hash))
  (pure [])

-- | A shard manager that uses base of stack repo hashes as shards,
-- and a dynamic shard assignment
shardByBaseOfStackRepoHash :: IO (Maybe [Text]) -> ShardManager Text
shardByBaseOfStackRepoHash refShardAssignment = ShardManager
  refShardAssignment
  (pure (\(BaseOfStack Glean.Repo{..}) _ -> repo_hash))
  (pure [])