git-annex/Remote/External/AsyncExtension.hs
2020-08-12 16:25:53 -04:00

125 lines
3.9 KiB
Haskell

{- External remote protocol async extension.
-
- Copyright 2013-2020 Joey Hess <id@joeyh.name>
-
- Licensed under the GNU AGPL version 3 or higher.
-}
{-# LANGUAGE BangPatterns #-}
module Remote.External.AsyncExtension where
import Common
import Remote.External.Types
import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Concurrent.STM.TBMChan
import qualified Data.Map.Strict as M
-- | Starts a thread that will handle all communication with the external
-- process. The input ExternalState communicates directly with the external
-- process.
runRelayToExternalAsync :: External -> ExternalState -> IO ExternalAsyncRelay
runRelayToExternalAsync external st = do
startcomm <- runRelayToExternalAsync' external st
return $ ExternalAsyncRelay $ do
(sendh, receiveh, shutdown) <- startcomm
return $ ExternalState
{ externalSend = atomically . writeTBMChan sendh
, externalReceive = fmap join $ atomically $ readTBMChan receiveh
-- This shuts down the whole relay.
, externalShutdown = shutdown
-- These three TVars are shared amoung all
-- ExternalStates that use this relay; they're
-- common state about the external process.
-- TODO: ALL code using these in Remote.External
-- has to be made async-safe.
, externalPrepared = externalPrepared st
, externalConfig = externalConfig st
, externalConfigChanges = externalConfigChanges st
}
runRelayToExternalAsync'
:: External
-> ExternalState
-> IO (IO (TBMChan String, TBMChan (Maybe String), Bool -> IO ()))
runRelayToExternalAsync' external st = do
newreqs <- newTVarIO []
startedcomms <- newTVarIO []
let startcomm = do
toq <- newTBMChanIO 10
fromq <- newTBMChanIO 10
let c = (toq, fromq, shutdown)
atomically $ do
l <- readTVar startedcomms
-- This append is ok because the maximum size
-- is the number of jobs that git-annex is
-- configured to use, which is a relatively
-- small number.
writeTVar startedcomms (l ++ [c])
return c
void $ async $ sendloop newreqs startedcomms
void $ async $ receiveloop newreqs M.empty
return startcomm
where
receiveloop newreqs jidmap = externalReceive st >>= \case
Just l -> case parseMessage l :: Maybe AsyncMessage of
Just (RESULT_ASYNC msg) -> getnext newreqs >>= \case
Just c -> do
relayto c msg
receiveloop newreqs jidmap
Nothing -> protoerr "unexpected RESULT-ASYNC"
Just (START_ASYNC jid msg) -> getnext newreqs >>= \case
Just c -> do
relayto c msg
let !jidmap' = M.insert jid c jidmap
receiveloop newreqs jidmap'
Nothing -> protoerr "unexpected START-ASYNC"
Just (END_ASYNC jid msg) -> case M.lookup jid jidmap of
Just c -> do
relayto c msg
closerelayto c
let !jidmap' = M.delete jid jidmap
receiveloop newreqs jidmap'
Nothing -> protoerr "END-ASYNC with unknown jobid"
Just (ASYNC jid msg) -> case M.lookup jid jidmap of
Just c -> do
relayto c msg
let !jidmap' = M.delete jid jidmap
receiveloop newreqs jidmap'
Nothing -> protoerr "ASYNC with unknown jobid"
_ -> protoerr "unexpected non-async message"
Nothing -> do
-- Unable to receive anything more from the
-- process, so it's not usable any longer.
forM_ (M.elems jidmap) closerelayto
shutdown True
sendloop newreqs startedcomms = do
error "TODO"
relayto (toq, _fromq) msg =
atomically $ writeTBMChan toq msg
closerelayto (toq, fromq) = do
atomically $ closeTBMChan toq
atomically $ closeTBMChan fromq
getnext l = atomically $ readTVar l >>= \case
[] -> return Nothing
(c:rest) -> do
writeTVar l rest
return (Just c)
shutdown b = do
r <- atomically $ do
r <- tryTakeTMVar (externalAsync external)
putTMVar (externalAsync external)
UncheckedExternalAsync
return r
case r of
Just (ExternalAsync _) -> externalShutdown st b
_ -> noop
protoerr s = giveup ("async special remote protocol error: " ++ s)