Add X-Eval-Alloc-Bytes response header to evaluation endpoints

Track GHC allocation bytes consumed per evaluation and return them
via response header for billing and resource monitoring.

Author: Thomas Gorissen

jl4-service/src/DataPlane.hs

@@ -37,7 +37,7 @@ import Data.Scientific (Scientific)
 import qualified Data.Set as Set
 import Data.Text (Text)
 import qualified Data.Text as Text
-import GHC.Conc (setAllocationCounter, enableAllocationLimit)
+import GHC.Conc (setAllocationCounter, getAllocationCounter, enableAllocationLimit)
 import GHC.IO.Exception (AllocationLimitExceeded (..))
 import Servant
 import System.FilePath ((<.>))
@@ -72,10 +72,14 @@ type DeploymentRoutes =
   :<|> "functions" :> Capture "name" Text :> FunctionRoutes
   :<|> "openapi.json" :> Get '[JSON] DeploymentMetadata  -- Simplified; returns metadata for now
 
+-- | Evaluation responses include an X-Eval-Alloc-Bytes header reporting
+-- the GHC allocation bytes consumed by the evaluation.
+type EvalResponse a = Headers '[Header "X-Eval-Alloc-Bytes" Int64] a
+
 type FunctionRoutes =
        Get '[JSON] Function
-  :<|> "evaluation" :> Header "X-L4-Trace" Text :> QueryParam "trace" TraceLevel :> QueryParam "graphviz" Bool :> ReqBody '[JSON] FnArguments :> Post '[JSON] SimpleResponse
-  :<|> "evaluation" :> "batch" :> Header "X-L4-Trace" Text :> QueryParam "trace" TraceLevel :> QueryParam "graphviz" Bool :> ReqBody '[JSON] BatchRequest :> Post '[JSON] BatchResponse
+  :<|> "evaluation" :> Header "X-L4-Trace" Text :> QueryParam "trace" TraceLevel :> QueryParam "graphviz" Bool :> ReqBody '[JSON] FnArguments :> Post '[JSON] (EvalResponse SimpleResponse)
+  :<|> "evaluation" :> "batch" :> Header "X-L4-Trace" Text :> QueryParam "trace" TraceLevel :> QueryParam "graphviz" Bool :> ReqBody '[JSON] BatchRequest :> Post '[JSON] (EvalResponse BatchResponse)
   :<|> "query-plan" :> ReqBody '[JSON] FnArguments :> Post '[JSON] QueryPlanResponse
   :<|> "state-graphs" :> Get '[JSON] StateGraphListResponse
   :<|> "state-graphs" :> Capture "graphName" Text :> Get '[PlainText] Text
@@ -163,7 +167,7 @@ evalFunctionHandler
   -> Maybe TraceLevel -- ?trace= query param
   -> Maybe Bool       -- ?graphviz= query param
   -> FnArguments
-  -> AppM SimpleResponse
+  -> AppM (EvalResponse SimpleResponse)
 evalFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz fnArgs = do
   (fns, _meta) <- requireDeploymentReady deployId
   vf <- requireFunction fns fnName
@@ -173,7 +177,7 @@ evalFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz fnArgs =
       isDeontic = Map.member "startTime" paramMap
                   && Map.member "events" paramMap
 
-  case (isDeontic, fnArgs.startTime, fnArgs.events) of
+  (result, allocBytes) <- case (isDeontic, fnArgs.startTime, fnArgs.events) of
     -- Non-deontic function: reject deontic params
     (False, Just _, _) ->
       throwError err400 { errBody = jsonError "startTime and events are only valid for functions returning DEONTIC" }
@@ -192,6 +196,7 @@ evalFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz fnArgs =
       runDeonticEvaluatorFor vf fnArgs.fnEvalBackend (Map.toList fnArgs.fnArguments) st evts
         vf.fnImpl.deonticPartyType vf.fnImpl.deonticActionType
         mTraceHeader mTraceParam mGraphViz
+  pure $ addHeader allocBytes result
 
 -- | POST /deployments/{id}/functions/{fn}/evaluation/batch
 batchFunctionHandler
@@ -201,7 +206,7 @@ batchFunctionHandler
   -> Maybe TraceLevel -- ?trace= query param
   -> Maybe Bool       -- ?graphviz= query param
   -> BatchRequest
-  -> AppM BatchResponse
+  -> AppM (EvalResponse BatchResponse)
 batchFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz batchArgs = do
   let traceLevel = determineTraceLevel mTraceHeader mTraceParam
       includeGraphViz = traceLevel == TraceFull && Maybe.fromMaybe False mGraphViz
@@ -211,7 +216,7 @@ batchFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz batchArg
   -- Capture the environment before going concurrent
   env <- ask
 
-  -- Evaluate all cases in parallel
+  -- Evaluate all cases in parallel, collecting alloc bytes per case
   evalResults <- liftIO $ forConcurrently batchArgs.cases $ \inputCase -> do
     let args = Map.assocs $ fmap Just inputCase.attributes
     r <- runAppM env (runEvaluatorForDirect vf Nothing args outputFilter traceLevel includeGraphViz)
@@ -223,12 +228,13 @@ batchFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz batchArg
     [] -> pure ()
 
   let
-    responses = [(rid, simpleResp) | (rid, Right simpleResp) <- evalResults]
+    responses = [(rid, simpleResp, alloc) | (rid, Right (simpleResp, alloc)) <- evalResults]
     nCases = length responses
+    totalAllocBytes = sum [alloc | (_, _, alloc) <- responses]
 
     successfulRuns =
       Maybe.mapMaybe
-        ( \(rid, simpleRes) -> case simpleRes of
+        ( \(rid, simpleRes, _) -> case simpleRes of
             SimpleResponse r -> Just (rid, r)
             SimpleError _ -> Nothing
         )
@@ -237,7 +243,7 @@ batchFunctionHandler deployId fnName mTraceHeader mTraceParam mGraphViz batchArg
     nSuccessful = length successfulRuns
     nIgnored = nCases - nSuccessful
 
-  pure $ BatchResponse
+  pure $ addHeader totalAllocBytes $ BatchResponse
     { cases =
         [ OutputCase
           { id = rid
@@ -341,7 +347,7 @@ openApiHandler deployId = do
 -- Evaluation helpers
 -- ----------------------------------------------------------------------------
 
--- | Run a function evaluator, returning SimpleResponse.
+-- | Run a function evaluator, returning SimpleResponse and alloc bytes.
 runEvaluatorFor
   :: ValidatedFunction
   -> Maybe EvalBackend
@@ -350,27 +356,27 @@ runEvaluatorFor
   -> Maybe Text       -- X-L4-Trace header
   -> Maybe TraceLevel -- ?trace= query param
   -> Maybe Bool       -- ?graphviz= query param
-  -> AppM SimpleResponse
+  -> AppM (SimpleResponse, Int64)
 runEvaluatorFor vf engine args outputFilter mTraceHeader mTraceParam mGraphViz = do
   let traceLevel = determineTraceLevel mTraceHeader mTraceParam
       includeGraphViz = traceLevel == TraceFull && Maybe.fromMaybe False mGraphViz
   runEvaluatorForDirect vf engine args outputFilter traceLevel includeGraphViz
 
--- | Core evaluator logic.
+-- | Core evaluator logic. Returns the response and GHC allocation bytes consumed.
 runEvaluatorForDirect
   :: ValidatedFunction
   -> Maybe EvalBackend
   -> [(Text, Maybe FnLiteral)]
   -> Maybe (Set.Set Text)
   -> TraceLevel
   -> Bool
-  -> AppM SimpleResponse
+  -> AppM (SimpleResponse, Int64)
 runEvaluatorForDirect vf engine args outputFilter traceLevel includeGraphViz = do
   let evalBackend = Maybe.fromMaybe JL4 engine
   case Map.lookup evalBackend vf.fnEvaluator of
     Nothing -> throwError err500 { errBody = jsonError "No evaluator available for backend" }
     Just runFn -> do
-      evaluationResult <-
+      (evaluationResult, allocBytes) <-
         timeoutAction $
           runExceptT
             ( runFn.runFunction
@@ -380,8 +386,8 @@ runEvaluatorForDirect vf engine args outputFilter traceLevel includeGraphViz = d
                 includeGraphViz
             )
       case evaluationResult of
-        Left err -> pure $ SimpleError err
-        Right r -> pure $ SimpleResponse r
+        Left err -> pure (SimpleError err, allocBytes)
+        Right r -> pure (SimpleResponse r, allocBytes)
 
 -- | Run deontic evaluation with EVALTRACE.
 runDeonticEvaluatorFor
@@ -395,7 +401,7 @@ runDeonticEvaluatorFor
   -> Maybe Text       -- X-L4-Trace header
   -> Maybe TraceLevel -- ?trace= query param
   -> Maybe Bool       -- ?graphviz= query param
-  -> AppM SimpleResponse
+  -> AppM (SimpleResponse, Int64)
 runDeonticEvaluatorFor vf _engine args startTime events mPartyType mActionType mTraceHeader mTraceParam mGraphViz = do
   let traceLevel = determineTraceLevel mTraceHeader mTraceParam
       includeGraphViz = traceLevel == TraceFull && Maybe.fromMaybe False mGraphViz
@@ -407,7 +413,7 @@ runDeonticEvaluatorFor vf _engine args startTime events mPartyType mActionType m
   let fnDecl = toDecl vf.fnImpl
       filepath = Text.unpack vf.fnImpl.name <.> "l4"
 
-  evaluationResult <-
+  (evaluationResult, allocBytes) <-
     timeoutAction $
       runExceptT
         ( evaluateWithCompiledDeontic
@@ -424,16 +430,17 @@ runDeonticEvaluatorFor vf _engine args startTime events mPartyType mActionType m
         )
 
   case evaluationResult of
-    Left err -> pure $ SimpleError err
-    Right r -> pure $ SimpleResponse r
+    Left err -> pure (SimpleError err, allocBytes)
+    Right r -> pure (SimpleResponse r, allocBytes)
 
 -- | Run an AppM action in IO with a given environment.
 runAppM :: AppEnv -> AppM a -> IO (Either ServerError a)
 runAppM env action = runHandler $ runReaderT action env
 
 -- | Timeout and memory-limited evaluation action.
 -- Uses configurable eval timeout and per-evaluation allocation limits.
-timeoutAction :: IO b -> AppM b
+-- Returns the result and the number of GHC allocation bytes consumed.
+timeoutAction :: IO b -> AppM (b, Int64)
 timeoutAction act = do
   cfg <- asks (.options)
   let timeoutMicros = cfg.evalTimeout * 1_000_000
@@ -442,12 +449,14 @@ timeoutAction act = do
     (timeout timeoutMicros $ do
       setAllocationCounter memLimitBytes
       enableAllocationLimit
-      act
+      r <- act
+      remaining <- getAllocationCounter
+      pure (r, memLimitBytes - remaining)
     ) `catch` \AllocationLimitExceeded ->
       pure Nothing
   case result of
     Nothing -> throwError err500 { errBody = jsonError "Evaluation resource limit exceeded" }
-    Just r -> pure r
+    Just x -> pure x
 
 -- ----------------------------------------------------------------------------
 -- Helpers