<commit_before>{-# LANGUAGE OverloadedStrings #-}
module MockServer
    (
      MockServer
    , mockServerHost
    , mockServerPort
    , withMockServer
    , recvMockServer
    , getConnectionCount
    ) where

import Data.Conduit
import Data.Conduit.Network
import Data.Conduit.Cereal
import Data.ByteString ( ByteString )
import Control.Concurrent
import Control.Monad.IO.Class
import Control.Monad.Catch (MonadThrow)
import Control.Exception
import Data.Serialize (Serialize, get)
import Data.MessagePack
import Data.Monoid
import qualified Data.Set as S

import Network.Fluent.Logger.Packable (pack)
import Network.Fluent.Logger.Unpackable (Unpackable, unpack)

data MockServer a = MockServer { mockServerChan :: Chan a
                               , mockServerThread :: ThreadId
                               , mockServerConnectionThreads :: MVar (S.Set ThreadId)
                               }

mockServerHost :: ByteString
mockServerHost = "127.0.0.1"

mockServerPort :: Int
mockServerPort = 24224

mockServerSettings :: ServerSettings
mockServerSettings = serverSettings mockServerPort "*"

app :: (MonadIO m, MonadThrow m, Serialize a) => Chan a -> MVar (S.Set ThreadId) -> AppData -> m ()
app chan mvconn ad = do
  myid <- liftIO myThreadId
  liftIO $ modifyMVar_ mvconn (return . S.insert myid)
  appSource ad $= conduitGet get $$ (addCleanup (unregister myid) $ sinkChan chan) where
    unregister myid _ = liftIO $ modifyMVar_ mvconn (return . S.delete myid)

sinkChan :: (MonadIO m, Serialize a) => Chan a -> Sink a m ()
sinkChan chan = do
  mx <- await
  case mx of
    Nothing -> return ()
    Just x  -> liftIO (writeChan chan x) >> sinkChan chan

withMockServer :: Serialize a => (MockServer a -> IO ()) -> IO ()
withMockServer = bracket runMockServer stopMockServer

recvMockServer :: Unpackable b => MockServer Object -> IO b
recvMockServer server = fmap unpack $ readChan (mockServerChan server)

runMockServer :: Serialize a => IO (MockServer a)
runMockServer = do
  chan <- newChan
  mvconn <- newMVar S.empty
  tid <- forkIO $ runTCPServer mockServerSettings $ app chan mvconn
  threadDelay 10000
  return MockServer { mockServerChan = chan
                    , mockServerThread = tid
                    , mockServerConnectionThreads = mvconn
                    }

stopMockServer :: Serialize a => MockServer a -> IO ()
stopMockServer server = do
  killThread $ mockServerThread server
  threadDelay 10000

getConnectionCount :: MockServer a -> IO Int
getConnectionCount s = fmap S.size $ readMVar $ mockServerConnectionThreads s<commit_msg>test/MockServer: kill all connection threads when stopping MockServer.<commit_after>{-# LANGUAGE OverloadedStrings #-}
module MockServer
    (
      MockServer
    , mockServerHost
    , mockServerPort
    , withMockServer
    , recvMockServer
    , getConnectionCount
    ) where

import Data.Conduit
import Data.Conduit.Network
import Data.Conduit.Cereal
import Data.ByteString ( ByteString )
import Control.Concurrent
import Control.Monad.IO.Class
import Control.Monad.Catch (MonadThrow)
import Control.Exception
import Data.Serialize (Serialize, get)
import Data.MessagePack
import Data.Monoid
import qualified Data.Set as S

import Network.Fluent.Logger.Packable (pack)
import Network.Fluent.Logger.Unpackable (Unpackable, unpack)

data MockServer a = MockServer { mockServerChan :: Chan a
                               , mockServerThread :: ThreadId
                               , mockServerConnectionThreads :: MVar (S.Set ThreadId)
                               }

mockServerHost :: ByteString
mockServerHost = "127.0.0.1"

mockServerPort :: Int
mockServerPort = 24224

mockServerSettings :: ServerSettings
mockServerSettings = serverSettings mockServerPort "*"

app :: (MonadIO m, MonadThrow m, Serialize a) => Chan a -> MVar (S.Set ThreadId) -> AppData -> m ()
app chan mvconn ad = do
  myid <- liftIO myThreadId
  liftIO $ modifyMVar_ mvconn (return . S.insert myid)
  appSource ad $= conduitGet get $$ (addCleanup (unregister myid) $ sinkChan chan) where
    unregister myid _ = liftIO $ modifyMVar_ mvconn (return . S.delete myid)

sinkChan :: (MonadIO m, Serialize a) => Chan a -> Sink a m ()
sinkChan chan = do
  mx <- await
  case mx of
    Nothing -> return ()
    Just x  -> liftIO (writeChan chan x) >> sinkChan chan

withMockServer :: Serialize a => (MockServer a -> IO ()) -> IO ()
withMockServer = bracket runMockServer stopMockServer

recvMockServer :: Unpackable b => MockServer Object -> IO b
recvMockServer server = fmap unpack $ readChan (mockServerChan server)

runMockServer :: Serialize a => IO (MockServer a)
runMockServer = do
  chan <- newChan
  mvconn <- newMVar S.empty
  tid <- forkIO $ runTCPServer mockServerSettings $ app chan mvconn
  threadDelay 10000
  return MockServer { mockServerChan = chan
                    , mockServerThread = tid
                    , mockServerConnectionThreads = mvconn
                    }

stopMockServer :: Serialize a => MockServer a -> IO ()
stopMockServer server = do
  killThread $ mockServerThread server
  (mapM_ killThread . S.toList) =<< (readMVar $ mockServerConnectionThreads server)
  threadDelay 10000

getConnectionCount :: MockServer a -> IO Int
getConnectionCount s = fmap S.size $ readMVar $ mockServerConnectionThreads s
{'fraction_non_alphanumeric': 0.08334908334908335, 'fraction_numerical': 0.007938007938007937, 'mean_word_length': 3.2643029814665594, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4695749', 'n_tokens_mistral': 1623, 'n_tokens_neox': 1530, 'n_words': 474}