<commit_before>module System.Mesos.Scheduler (
  Scheduler,
  ToScheduler(..),
  createScheduler,
  destroyScheduler,
  SchedulerDriver,
  create,
  destroyDriver,
  start,
  stop,
  abort,
  joinDriver,
  run,
  requestResources,
  launchTasks,
  killTask,
  declineOffer,
  reviveOffers,
  sendFrameworkMessage,
  reconcileTasks
) where
import Data.ByteString (ByteString, packCStringLen)
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Array
import System.Mesos.Internal

class ToScheduler a where
  registered :: a -> SchedulerDriver -> FrameworkID -> MasterInfo -> IO ()
  registered _ _ _ _ = return ()

  reRegistered :: a -> SchedulerDriver -> MasterInfo -> IO ()
  reRegistered _ _ _ = return ()

  disconnected :: a -> SchedulerDriver -> IO ()
  disconnected _ _ = return ()

  resourceOffers :: a -> SchedulerDriver -> [Offer] -> IO ()
  resourceOffers _ _ _ = return ()

  offerRescinded :: a -> SchedulerDriver -> OfferID -> IO ()
  offerRescinded _ _ _ = return ()

  statusUpdate :: a -> SchedulerDriver -> TaskStatus -> IO ()
  statusUpdate _ _ _ = return ()

  frameworkMessage :: a -> SchedulerDriver -> ExecutorID -> SlaveID -> ByteString -> IO ()
  frameworkMessage  _ _ _ _ _ = return ()

  slaveLost :: a -> SchedulerDriver -> SlaveID -> IO ()
  slaveLost _ _ _ = return ()

  executorLost :: a -> SchedulerDriver -> ExecutorID -> SlaveID -> Status -> IO ()
  executorLost _ _ _ _ _ = return ()

  errorMessage :: a -> SchedulerDriver -> ByteString -> IO ()
  errorMessage _ _ _ = return ()

createScheduler :: ToScheduler a => a -> IO Scheduler
createScheduler s = do
  registeredFun <- wrapSchedulerRegistered $ \sdp fp mp -> do
    let sd = SchedulerDriver sdp
    f <- unmarshal fp
    m <- unmarshal mp
    (registered s) sd f m
  reRegisteredFun <- wrapSchedulerReRegistered $ \sdp mip -> do
    let sd = SchedulerDriver sdp
    mi <- unmarshal mip
    (reRegistered s) sd mi
  disconnectedFun <- wrapSchedulerDisconnected $ \sdp -> do
    let sd = SchedulerDriver sdp
    (disconnected s) sd
  resourceOffersFun <- wrapSchedulerResourceOffers $ \sdp os c -> do
    let sd = SchedulerDriver sdp
    offers <- mapM unmarshal =<< peekArray (fromIntegral c) os
    (resourceOffers s) sd offers
  offerRescindedFun <- wrapSchedulerOfferRescinded $ \sdp oidp -> do
    let sd = SchedulerDriver sdp
    oid <- unmarshal oidp
    (offerRescinded s) sd oid
  statusUpdateFun <- wrapSchedulerStatusUpdate $ \sdp tsp -> do
    let sd = SchedulerDriver sdp
    ts <- unmarshal tsp
    (statusUpdate s) sd ts
  frameworkMessageFun <- wrapSchedulerFrameworkMessage $ \sdp eip sip ptr c -> do
    let sd = SchedulerDriver sdp
    ei <- unmarshal eip
    si <- unmarshal sip
    bs <- packCStringLen (ptr, fromIntegral c)
    (frameworkMessage s) sd ei si bs
  slaveLostFun <- wrapSchedulerSlaveLost $ \sdp sip -> do
    let sd = SchedulerDriver sdp
    si <- unmarshal sip
    (slaveLost s) sd si
  executorLostFun <- wrapSchedulerExecutorLost $ \sdp eip sip st -> do
    let sd = SchedulerDriver sdp
    ei <- unmarshal eip
    si <- unmarshal sip
    (executorLost s) sd ei si (toEnum $ fromIntegral st)
  errorFun <- wrapSchedulerError $ \sdp ptr c -> do
    let sd = SchedulerDriver sdp
    bs <- packCStringLen (ptr, fromIntegral c)
    (errorMessage s) sd bs
  schedulerPtr <- c_createScheduler registeredFun reRegisteredFun disconnectedFun resourceOffersFun offerRescindedFun statusUpdateFun frameworkMessageFun slaveLostFun executorLostFun errorFun
  return $ Scheduler schedulerPtr registeredFun reRegisteredFun disconnectedFun resourceOffersFun offerRescindedFun statusUpdateFun frameworkMessageFun slaveLostFun executorLostFun errorFun

destroyScheduler :: Scheduler -> IO ()
destroyScheduler s = do
  c_destroyScheduler $ schedulerImpl s
  freeHaskellFunPtr $ rawSchedulerRegistered s
  freeHaskellFunPtr $ rawSchedulerReRegistered s
  freeHaskellFunPtr $ rawSchedulerDisconnected s
  freeHaskellFunPtr $ rawSchedulerResourceOffers s
  freeHaskellFunPtr $ rawSchedulerOfferRescinded s
  freeHaskellFunPtr $ rawSchedulerStatusUpdate s
  freeHaskellFunPtr $ rawSchedulerFrameworkMessage s
  freeHaskellFunPtr $ rawSchedulerSlaveLost s
  freeHaskellFunPtr $ rawSchedulerExecutorLost s
  freeHaskellFunPtr $ rawSchedulerError s

excerciseMethods :: Scheduler -> IO ()
excerciseMethods = c_exerciseMethods . schedulerImpl

withDriver :: (SchedulerDriverPtr -> IO CInt) -> SchedulerDriver -> IO Status
withDriver f (SchedulerDriver p) = fmap (toEnum . fromIntegral) $ f p

create :: Scheduler -> FrameworkInfo -> ByteString -> Maybe Credential -> IO SchedulerDriver
create = undefined

destroyDriver :: SchedulerDriver -> IO ()
destroyDriver = c_destroySchedulerDriver . fromSchedulerDriver

start :: SchedulerDriver -> IO Status
start = withDriver c_startSchedulerDriver

stop :: SchedulerDriver -> Bool -> IO Status
stop d f = withDriver (\p -> c_stopSchedulerDriver p fi) d
  where
    fi = if f then 1 else 0

abort :: SchedulerDriver -> IO Status
abort = withDriver c_abortSchedulerDriver

joinDriver :: SchedulerDriver -> IO Status
joinDriver = withDriver c_joinSchedulerDriver
<commit_msg>Simplify resource management for scheduler driver<commit_after>module System.Mesos.Scheduler (
  Scheduler,
  ToScheduler(..),
  createScheduler,
  destroyScheduler,
  withSchedulerDriver,
  SchedulerDriver,
  create,
  destroyDriver,
  start,
  stop,
  abort,
  joinDriver,
  run,
  requestResources,
  launchTasks,
  killTask,
  declineOffer,
  reviveOffers,
  sendFrameworkMessage,
  reconcileTasks
) where
import Data.ByteString (ByteString, packCStringLen)
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Array
import System.Mesos.Internal

class ToScheduler a where
  registered :: a -> SchedulerDriver -> FrameworkID -> MasterInfo -> IO ()
  registered _ _ _ _ = return ()

  reRegistered :: a -> SchedulerDriver -> MasterInfo -> IO ()
  reRegistered _ _ _ = return ()

  disconnected :: a -> SchedulerDriver -> IO ()
  disconnected _ _ = return ()

  resourceOffers :: a -> SchedulerDriver -> [Offer] -> IO ()
  resourceOffers _ _ _ = return ()

  offerRescinded :: a -> SchedulerDriver -> OfferID -> IO ()
  offerRescinded _ _ _ = return ()

  statusUpdate :: a -> SchedulerDriver -> TaskStatus -> IO ()
  statusUpdate _ _ _ = return ()

  frameworkMessage :: a -> SchedulerDriver -> ExecutorID -> SlaveID -> ByteString -> IO ()
  frameworkMessage  _ _ _ _ _ = return ()

  slaveLost :: a -> SchedulerDriver -> SlaveID -> IO ()
  slaveLost _ _ _ = return ()

  executorLost :: a -> SchedulerDriver -> ExecutorID -> SlaveID -> Status -> IO ()
  executorLost _ _ _ _ _ = return ()

  errorMessage :: a -> SchedulerDriver -> ByteString -> IO ()
  errorMessage _ _ _ = return ()

createScheduler :: ToScheduler a => a -> IO Scheduler
createScheduler s = do
  registeredFun <- wrapSchedulerRegistered $ \sdp fp mp -> do
    let sd = SchedulerDriver sdp
    f <- unmarshal fp
    m <- unmarshal mp
    (registered s) sd f m
  reRegisteredFun <- wrapSchedulerReRegistered $ \sdp mip -> do
    let sd = SchedulerDriver sdp
    mi <- unmarshal mip
    (reRegistered s) sd mi
  disconnectedFun <- wrapSchedulerDisconnected $ \sdp -> do
    let sd = SchedulerDriver sdp
    (disconnected s) sd
  resourceOffersFun <- wrapSchedulerResourceOffers $ \sdp os c -> do
    let sd = SchedulerDriver sdp
    offers <- mapM unmarshal =<< peekArray (fromIntegral c) os
    (resourceOffers s) sd offers
  offerRescindedFun <- wrapSchedulerOfferRescinded $ \sdp oidp -> do
    let sd = SchedulerDriver sdp
    oid <- unmarshal oidp
    (offerRescinded s) sd oid
  statusUpdateFun <- wrapSchedulerStatusUpdate $ \sdp tsp -> do
    let sd = SchedulerDriver sdp
    ts <- unmarshal tsp
    (statusUpdate s) sd ts
  frameworkMessageFun <- wrapSchedulerFrameworkMessage $ \sdp eip sip ptr c -> do
    let sd = SchedulerDriver sdp
    ei <- unmarshal eip
    si <- unmarshal sip
    bs <- packCStringLen (ptr, fromIntegral c)
    (frameworkMessage s) sd ei si bs
  slaveLostFun <- wrapSchedulerSlaveLost $ \sdp sip -> do
    let sd = SchedulerDriver sdp
    si <- unmarshal sip
    (slaveLost s) sd si
  executorLostFun <- wrapSchedulerExecutorLost $ \sdp eip sip st -> do
    let sd = SchedulerDriver sdp
    ei <- unmarshal eip
    si <- unmarshal sip
    (executorLost s) sd ei si (toEnum $ fromIntegral st)
  errorFun <- wrapSchedulerError $ \sdp ptr c -> do
    let sd = SchedulerDriver sdp
    bs <- packCStringLen (ptr, fromIntegral c)
    (errorMessage s) sd bs
  schedulerPtr <- c_createScheduler registeredFun reRegisteredFun disconnectedFun resourceOffersFun offerRescindedFun statusUpdateFun frameworkMessageFun slaveLostFun executorLostFun errorFun
  return $ Scheduler schedulerPtr registeredFun reRegisteredFun disconnectedFun resourceOffersFun offerRescindedFun statusUpdateFun frameworkMessageFun slaveLostFun executorLostFun errorFun

destroyScheduler :: Scheduler -> IO ()
destroyScheduler s = do
  c_destroyScheduler $ schedulerImpl s
  freeHaskellFunPtr $ rawSchedulerRegistered s
  freeHaskellFunPtr $ rawSchedulerReRegistered s
  freeHaskellFunPtr $ rawSchedulerDisconnected s
  freeHaskellFunPtr $ rawSchedulerResourceOffers s
  freeHaskellFunPtr $ rawSchedulerOfferRescinded s
  freeHaskellFunPtr $ rawSchedulerStatusUpdate s
  freeHaskellFunPtr $ rawSchedulerFrameworkMessage s
  freeHaskellFunPtr $ rawSchedulerSlaveLost s
  freeHaskellFunPtr $ rawSchedulerExecutorLost s
  freeHaskellFunPtr $ rawSchedulerError s

excerciseMethods :: Scheduler -> IO ()
excerciseMethods = c_exerciseMethods . schedulerImpl

withDriver :: (SchedulerDriverPtr -> IO CInt) -> SchedulerDriver -> IO Status
withDriver f (SchedulerDriver p) = fmap (toEnum . fromIntegral) $ f p

withSchedulerDriver :: ToScheduler a => a -> FrameworkInfo -> ByteString -> Maybe Credential -> (SchedulerDriver -> IO b) -> IO b
withSchedulerDriver s i h c f = do
  scheduler <- createScheduler s
  driver <- create scheduler i h c
  result <- f driver
  destroyDriver driver
  destroyScheduler scheduler
  return result

create :: Scheduler -> FrameworkInfo -> ByteString -> Maybe Credential -> IO SchedulerDriver
create s i h mc = do
  fiP <- marshal i
  result <- unsafeUseAsCStringLen h $ \(hP, hLen) -> case mc of
    Nothing -> c_createSchedulerDriver (schedulerImpl s) fiP hP (fromIntegral hLen)
    Just c -> do
      cP <- marshal c
      result <- c_createSchedulerDriverWithCredentials (schedulerImpl s) fiP hP (fromIntegral hLen) cP
      destroy cP
      return result
  destroy fiP
  return $ SchedulerDriver result

destroyDriver :: SchedulerDriver -> IO ()
destroyDriver = c_destroySchedulerDriver . fromSchedulerDriver

start :: SchedulerDriver -> IO Status
start = withDriver c_startSchedulerDriver

stop :: SchedulerDriver -> Bool -> IO Status
stop d f = withDriver (\p -> c_stopSchedulerDriver p fi) d
  where
    fi = if f then 1 else 0

abort :: SchedulerDriver -> IO Status
abort = withDriver c_abortSchedulerDriver

joinDriver :: SchedulerDriver -> IO Status
joinDriver = withDriver c_joinSchedulerDriver

{'fraction_non_alphanumeric': 0.07899461400359066, 'fraction_numerical': 0.0003590664272890485, 'mean_word_length': 3.759077317385733, 'pattern_counts': {'":': 0, '<': 62, '<?xml version=': 0, '>': 137, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3283955', 'n_tokens_mistral': 3508, 'n_tokens_neox': 3321, 'n_words': 1009}