flink-statefulfun 0.3.0.0 → 0.4.0.0

raw patch · 7 files changed

+348/−313 lines, 7 filesdep +microlensdep +timedep +uuidPVP ok

version bump matches the API change (PVP)

Dependencies added: microlens, time, uuid

API changes (from Hackage documentation)

Files

• README.md +12/−26
• flink-statefulfun.cabal +9/−2
• proto/Example.proto +0/−11
• proto/Kafka.proto +11/−10
• proto/RequestReply.proto +161/−126
• src/Network/Flink/Internal/Stateful.hs +141/−129
• src/Network/Flink/Stateful.hs +14/−9

@@ -82,7 +82,7 @@   newCount <- (+ 1) <$> insideCtx greeterStateCount   let respMsg = "Saw " <> T.unpack name <> " " <> show newCount <> " time(s)" -  sendMsg ("printing", "printer") (response $ T.pack respMsg)+  sendMsgProto ("printing", "printer") (response $ T.pack respMsg)   modifyCtx (\old -> old {greeterStateCount = newCount})   where     name = msg ^. EX.name@@ -115,17 +115,15 @@   putStrLn "http://localhost:8000/"   run 8000 (logStdout $ createApp functionTable) -greeterState = BSL.toStrict . encode $ GreeterState 0- functionTable :: FunctionTable functionTable =   Map.fromList-    [ (("printing", "printer"), ("", flinkWrapper . protoInput $ printer)),-      (("greeting", "counter"), (greeterState, flinkWrapper . protoInput . jsonState $ counter))+    [ ((FuncType "greeting" "greeterEntry"), flinkWrapper () (Expiration NONE 0) (greeterEntry . getProto)),+      ((FuncType  "greeting" "counter"), flinkWrapper (JsonSerde (GreeterState 0)) (Expiration AFTER_CALL 5) (jsonState . counter . getProto))     ] ``` -The `FunctionTable` is a Map of `(namespace, functionType)` to `(initialState, wrappedFunction)`.+The `FunctionTable` is a Map of `(namespace, functionType)` to `wrappedFunction`. `jsonState` is a helper to serialize your function state as `JSON` for storage in the flink backend. `protoState` can also be used if that is your preference. `flinkWrapper` transforms your function into one that takes arbitrary `ByteString`s so that every function in the@@ -143,30 +141,18 @@ To use the functions that are now served from the API, we now need to declare it in the `module.yaml`.  ```yaml-version: "1.0"+version: "3.0" module:   meta:     type: remote   spec:-    functions:-      - function:-          meta:-            kind: http-            type: greeting/counter-          spec:-            endpoint: http://localhost:8000/statefun-            states:-              - flink_state-            maxNumBatchRequests: 500-            timeout: 2min-      - function:-          meta:-            kind: http-            type: printing/printer-          spec:-            endpoint: http://localhost:8000/statefun-            maxNumBatchRequests: 500-            timeout: 2min+    endpoints:+      - endpoint:+        meta: +          kind: http+        spec:+          functions: greeting/*+          urlPathTemplate: http://localhost:8000/statefun ```  Flink Statefun supports multiple states, but for simplicity the SDK just serializes the whole

@@ -1,7 +1,7 @@ cabal-version:       2.4  name:                flink-statefulfun-version:             0.3.0.0+version:             0.4.0.0 synopsis:            Flink stateful functions SDK description:     Typeclasses for serving Flink stateful functions@@ -20,7 +20,8 @@                      README.md                      proto/**/*.proto -tested-with: GHC==8.10.2,+tested-with: GHC==8.10.5,+             GHC==8.10.2,              GHC==8.8.4,              GHC==8.8.3,              GHC==8.6.5@@ -55,6 +56,7 @@                      , warp >= 3 && < 3.4                      , http-types >= 0.10 && < 0.13                      , http-media >= 0.7 && < 0.9+                     , microlens == 0.4.12.0                      , mtl >= 2 && < 2.3                      , lens-family >= 2 && < 2.2                      , proto-lens >= 0.5 && < 0.8@@ -62,6 +64,8 @@                      , proto-lens-protobuf-types >= 0.5 && < 0.8                      , servant >= 0.16 && < 0.19                      , servant-server >= 0.16 && < 0.19+                     , time == 1.9.3+                     , uuid == 1.3.15    default-extensions:  OverloadedStrings                        QuasiQuotes@@ -79,5 +83,8 @@                        MultiParamTypeClasses                        FunctionalDependencies                        ConstraintKinds+                       ScopedTypeVariables+                       TypeApplications+                       RankNTypes    default-language:    Haskell2010

@@ -1,11 +0,0 @@-syntax = "proto3";--package example;--message GreeterRequest {-  string name = 1;-}--message GreeterResponse {-  string greeting = 1;-}

@@ -16,14 +16,15 @@  * limitations under the License.  */ -syntax = "proto3";--package org.apache.flink.statefun.flink.io;-option java_package = "org.apache.flink.statefun.flink.io.generated";-option java_multiple_files = true;+ syntax = "proto3"; -message KafkaProducerRecord {-    string key = 1;-    bytes value_bytes = 2;-    string topic = 3;-}+ package io.statefun.sdk.egress;+ + option java_package = "org.apache.flink.statefun.sdk.egress.generated";+ option java_multiple_files = true;+ + message KafkaProducerRecord {+     string key = 1;+     bytes value_bytes = 2;+     string topic = 3;+ }

@@ -16,130 +16,165 @@  * limitations under the License.  */ -syntax = "proto3";--package org.apache.flink.statefun.flink.core.polyglot;-option java_package = "org.apache.flink.statefun.flink.core.polyglot.generated";-option java_multiple_files = true;--import "google/protobuf/any.proto";--// --------------------------------------------------------------------------------------------------------------------// Common message definitions-// ---------------------------------------------------------------------------------------------------------------------// An Address is the unique identity of an individual StatefulFunction, containing-// a function's type and an unique identifier within the type. The function's-// type denotes the "class" of function to invoke, while the unique identifier addresses the-// invocation to a specific function instance.-message Address {-    string namespace = 1;-    string type = 2;-    string id = 3;-}--// --------------------------------------------------------------------------------------------------------------------// Messages sent to a Remote Function  -// ---------------------------------------------------------------------------------------------------------------------// The following section contains all the message types that are sent -// from Flink to a remote function.-message ToFunction {-    // PersistedValue represents a PersistedValue's value that is managed by Flink on behalf of a remote function. -    message PersistedValue {-        // The unique name of the persisted state.-        string state_name = 1;-        // The serialized state value-        bytes state_value = 2;-    }--    // Invocation represents a remote function call, it associated with an (optional) return address,-    // and an argument. -    message Invocation {-        // The address of the function that requested the invocation (possibly absent)-        Address caller = 1;-        // The invocation argument (aka the message sent to the target function)-        google.protobuf.Any argument = 2;-    }--    // InvocationBatchRequest represents a request to invoke a remote function. It is always associated with a target-    // address (the function to invoke), a list of eager state values.-    message InvocationBatchRequest {-        // The address of the function to invoke-        Address target = 1;-        // A list of PersistedValues that were registered as an eager state.-        repeated PersistedValue state = 2;-        // A non empty (at least one) list of invocations-        repeated Invocation invocations = 3;-    }--    oneof request {-        InvocationBatchRequest invocation = 100;-    }-}--// --------------------------------------------------------------------------------------------------------------------// Messages sent from a Remote Function  -// ---------------------------------------------------------------------------------------------------------------------// The following section contains messages sent from a remote function back to Flink. -message FromFunction {-    // MutatePersistedValueCommand represents a command sent from a remote function to Flink,-    // requesting a change to a persisted value.-    message PersistedValueMutation {-        enum MutationType {-            DELETE = 0;-            MODIFY = 1;-        }-        MutationType mutation_type = 1;-        string state_name = 2;-        bytes state_value = 3;-    }--    // Invocation represents a remote function call, it associated with a (mandatory) target address,-    // and an argument. -    message Invocation {-        // The target function to invoke -        Address target = 1;-        // The invocation argument (aka the message sent to the target function)-        google.protobuf.Any argument = 2;-    }--    // DelayedInvocation represents a delayed remote function call with a target address, an argument-    // and a delay in milliseconds, after which this message to be sent.-    message DelayedInvocation {-        // the amount of milliseconds to wait before sending this message-        int64 delay_in_ms = 1;-        // the target address to send this message to-        Address target = 2;-        // the invocation argument-        google.protobuf.Any argument = 3;-    }--    // EgressMessage an argument to forward to an egress.-    // An egress is identified by a namespace and type (see EgressIdentifier SDK class).-    // The argument is a google.protobuf.Any-    message EgressMessage {-        // The target egress namespace-        string egress_namespace = 1;-        // The target egress type-        string egress_type = 2;-        // egress argument-        google.protobuf.Any argument = 3;-    }--    // InvocationResponse represents a result of an org.apache.flink.statefun.flink.core.polyglot.ToFunction.InvocationBatchRequest-    // it contains a list of state mutation to preform as a result of computing this batch, and a list of outgoing messages.-    message InvocationResponse {-        repeated PersistedValueMutation state_mutations = 1;-        repeated Invocation outgoing_messages = 2;-        repeated DelayedInvocation delayed_invocations = 3;-        repeated EgressMessage outgoing_egresses = 4;-    }--    oneof response {-        InvocationResponse invocation_result = 100;-    }-}-+ syntax = "proto3"; + package io.statefun.sdk.reqreply;+ + option java_package = "org.apache.flink.statefun.sdk.reqreply.generated";+ option java_multiple_files = true;+ + // -------------------------------------------------------------------------------------------------------------------+ // Common message definitions+ // -------------------------------------------------------------------------------------------------------------------+ + // An Address is the unique identity of an individual StatefulFunction, containing+ // a function's type and an unique identifier within the type. The function's+ // type denotes the "class" of function to invoke, while the unique identifier addresses the+ // invocation to a specific function instance.+ message Address {+     string namespace = 1;+     string type = 2;+     string id = 3;+ }+ + message TypedValue {+     string typename = 1;+     // has_value is set to differentiate a zero length value bytes explicitly set,+     // or a non existing value.+     bool has_value = 2;+     bytes value = 3;+ }+ + // -------------------------------------------------------------------------------------------------------------------+ // Messages sent to a Remote Function  + // -------------------------------------------------------------------------------------------------------------------+ + // The following section contains all the message types that are sent + // from Flink to a remote function.+ message ToFunction {+     // PersistedValue represents a PersistedValue's value that is managed by Flink on behalf of a remote function. +     message PersistedValue {+         // The unique name of the persisted state.+         string state_name = 1;+         // The serialized state value+         TypedValue state_value = 2;+     }+ +     // Invocation represents a remote function call, it associated with an (optional) return address,+     // and an argument. +     message Invocation {+         // The address of the function that requested the invocation (possibly absent)+         Address caller = 1;+         // The invocation argument (aka the message sent to the target function)+         TypedValue argument = 2;+     }+ +     // InvocationBatchRequest represents a request to invoke a remote function. It is always associated with a target+     // address (the function to invoke), and a list of values for registered state.+     message InvocationBatchRequest {+         // The address of the function to invoke+         Address target = 1;+         // A list of PersistedValues that were registered as a persisted state.+         repeated PersistedValue state = 2;+         // A non empty (at least one) list of invocations+         repeated Invocation invocations = 3;+     }+ +     oneof request {+         InvocationBatchRequest invocation = 100;+     }+ }+ + // -------------------------------------------------------------------------------------------------------------------+ // Messages sent from a Remote Function  + // -------------------------------------------------------------------------------------------------------------------+ + // The following section contains messages sent from a remote function back to Flink. + message FromFunction {+     // MutatePersistedValueCommand represents a command sent from a remote function to Flink,+     // requesting a change to a persisted value.+     message PersistedValueMutation {+         enum MutationType {+             DELETE = 0;+             MODIFY = 1;+         }+         MutationType mutation_type = 1;+         string state_name = 2;+         TypedValue state_value = 3;+     }+ +     // Invocation represents a remote function call, it associated with a (mandatory) target address,+     // and an argument. +     message Invocation {+         // The target function to invoke +         Address target = 1;+         // The invocation argument (aka the message sent to the target function)+         TypedValue argument = 2;+     }+ +     // DelayedInvocation represents a delayed remote function call with a target address, an argument+     // and a delay in milliseconds, after which this message to be sent.+     message DelayedInvocation {+         // the amount of milliseconds to wait before sending this message+         int64 delay_in_ms = 1;+         // the target address to send this message to+         Address target = 2;+         // the invocation argument+         TypedValue argument = 3;+     }+ +     // EgressMessage an argument to forward to an egress.+     // An egress is identified by a namespace and type (see EgressIdentifier SDK class).+     // The argument is an io.statefun.sdk.reqreply.TypedValue.+     message EgressMessage {+         // The target egress namespace+         string egress_namespace = 1;+         // The target egress type+         string egress_type = 2;+         // egress argument+         TypedValue argument = 3;+     }+ +     // InvocationResponse represents a result of an io.statefun.sdk.reqreply.ToFunction.InvocationBatchRequest+     // it contains a list of state mutation to preform as a result of computing this batch, and a list of outgoing messages.+     message InvocationResponse {+         repeated PersistedValueMutation state_mutations = 1;+         repeated Invocation outgoing_messages = 2;+         repeated DelayedInvocation delayed_invocations = 3;+         repeated EgressMessage outgoing_egresses = 4;+     }+ +     // ExpirationSpec represents TTL (Time-To-Live) configuration for persisted states.+     message ExpirationSpec {+         enum ExpireMode {+             NONE = 0;+             AFTER_WRITE = 1;+             AFTER_INVOKE = 2;+         }+         ExpireMode mode = 1;+         int64 expire_after_millis = 2;+     }+ +     // PersistedValueSpec represents specifications of a function's persisted value state.+     message PersistedValueSpec {+         string state_name = 1;+         ExpirationSpec expiration_spec = 2;+         string type_typename = 3;+     }+ +     // IncompleteInvocationContext represents a result of an io.statefun.sdk.reqreply.ToFunction.InvocationBatchRequest,+     // which should be used as the response if the InvocationBatchRequest provided incomplete information about the+     // invocation, e.g. insufficient state values were provided.+     message IncompleteInvocationContext {+         repeated PersistedValueSpec missing_values = 1;+     }+ +     // Response sent from the function, as a result of an io.statefun.sdk.reqreply.ToFunction.InvocationBatchRequest.+     // It can be one of the following types:+     //   - io.statefun.sdk.reqreply.FromFunction.InvocationResponse+     //   - io.statefun.sdk.reqreply.FromFunction.IncompleteInvocationContext+     oneof response {+         InvocationResponse invocation_result = 100;+         IncompleteInvocationContext incomplete_invocation_context = 101;+     }+ } 

@@ -1,3 +1,4 @@+{-# LANGUAGE PatternSynonyms #-} module Network.Flink.Internal.Stateful   ( StatefulFunc       ( insideCtx,@@ -6,28 +7,29 @@         modifyCtx,         sendMsg,         sendMsgDelay,-        sendEgressMsg,-        sendByteMsg,-        sendByteMsgDelay+        sendEgressMsg       ),     flinkWrapper,     createApp,     flinkServer,     flinkApi,+    Address(.., Address'),+    FuncType (..),     Function (..),     Serde (..),     FunctionState (..),     FlinkError (..),     FunctionTable,     Env (..),+    Expiration(..),+    ExpirationMode(..),     newState,     ProtoSerde (..),     JsonSerde (..),     jsonState,     protoState,-    serdeInput,-    protoInput,-    jsonInput,+    sendProtoMsg,+    sendProtoMsgDelay   ) where @@ -41,29 +43,38 @@ import Data.Foldable (Foldable (toList)) import Data.Map (Map) import qualified Data.Map as Map-import Data.Maybe (listToMaybe)-import Data.ProtoLens (Message, defMessage, encodeMessage)+import Data.ProtoLens (Message, defMessage, encodeMessage, messageName) import Data.ProtoLens.Any (UnpackError)-import qualified Data.ProtoLens.Any as Any import Data.ProtoLens.Encoding (decodeMessage) import Data.ProtoLens.Prism import Data.Sequence (Seq) import qualified Data.Sequence as Seq import Data.Text (Text) import Data.Text.Lazy (fromStrict)+import Data.Coerce ( coerce ) import qualified Data.Text.Lazy.Encoding as T import Lens.Family2+import Lens.Micro ( traversed, filtered ) import Network.Flink.Internal.ProtoServant (Proto)-import Proto.Google.Protobuf.Any (Any)-import qualified Proto.Google.Protobuf.Any_Fields as Any import Proto.RequestReply (FromFunction, ToFunction) import qualified Proto.RequestReply as PR import qualified Proto.RequestReply_Fields as PR import Servant ---- | Table of stateful functions `(functionNamespace, functionType) -> (initialState, function)-type FunctionTable = Map (Text, Text) (ByteString, ByteString -> Env -> PR.ToFunction'InvocationBatchRequest -> IO (Either FlinkError (FunctionState ByteString)))+import Data.Time.Clock ( NominalDiffTime ) +data FuncType = FuncType Text Text deriving (Eq, Ord)+data Address = Address FuncType Text++{-# COMPLETE Address' #-}+pattern Address' :: Text -> Text -> Text -> Address+pattern Address' fnamespace fnTp fid = Address (FuncType fnamespace fnTp) fid++data FuncRes = IncompleteContext Expiration Text | UpdatedState (FunctionState PR.TypedValue) deriving Show+type FuncExec = Env -> PR.ToFunction'InvocationBatchRequest -> IO (Either FlinkError FuncRes)+--- | Table of stateful functions `(functionNamespace, functionType) -> function+type FunctionTable = Map FuncType FuncExec+ data Env = Env   { envFunctionNamespace :: Text,     envFunctionType :: Text,@@ -83,22 +94,35 @@ newState :: a -> FunctionState a newState initialCtx = FunctionState initialCtx False mempty mempty mempty +data ExpirationMode =  NONE | AFTER_WRITE | AFTER_CALL deriving (Show, Eq)++data Expiration = Expiration {+  emode :: ExpirationMode,+  expireAfterMillis :: NominalDiffTime+} deriving (Show, Eq)+ -- | Monad stack used for the execution of a Flink stateful function -- Don't reference this directly in your code if possible newtype Function s a = Function {runFunction :: ExceptT FlinkError (StateT (FunctionState s) (ReaderT Env IO)) a}   deriving (Monad, Applicative, Functor, MonadState (FunctionState s), MonadError FlinkError, MonadIO, MonadReader Env)  class Serde a where+  -- | Type name+  tpName :: Proxy a -> Text+   -- | decodes types from strict 'ByteString's   deserializeBytes :: ByteString -> Either String a    -- | encodes types to strict 'ByteString's   serializeBytes :: a -> ByteString -newtype ProtoSerde a = ProtoSerde {getMessage :: a}+newtype ProtoSerde a = ProtoSerde {getProto :: a}   deriving (Functor)  instance Message a => Serde (ProtoSerde a) where+  tpName px = "type.googleapis.com/" <> messageName (unliftP px)+    where unliftP :: Proxy (f a) -> Proxy a+          unliftP Proxy = Proxy   deserializeBytes a = ProtoSerde <$> decodeMessage a   serializeBytes (ProtoSerde a) = encodeMessage a @@ -108,18 +132,22 @@   deriving (Functor)  instance Json a => Serde (JsonSerde a) where+  tpName _ = "json/json" -- TODO: add adt name   deserializeBytes a = JsonSerde <$> eitherDecode (BSL.fromStrict a)   serializeBytes (JsonSerde a) = BSL.toStrict $ encode a  instance Serde () where+  tpName _ = "ghc/Unit"   deserializeBytes _ = pure ()   serializeBytes _ = ""  instance Serde ByteString where+  tpName _ = "ghc/Data.ByteString"   deserializeBytes = pure   serializeBytes = id  instance Serde BSL.ByteString where+  tpName _ = "ghc/Data.ByteString.Lazy"   deserializeBytes = pure . BSL.fromStrict   serializeBytes = BSL.toStrict @@ -140,22 +168,6 @@   getCtx :: m s   setCtx :: s -> m ()   modifyCtx :: (s -> s) -> m ()-  sendMsg ::-    Message a =>-    -- | Function address (namespace, type, id)-    (Text, Text, Text) ->-    -- | protobuf message to send-    a ->-    m ()-  sendMsgDelay ::-    Message a =>-    -- | Function address (namespace, type, id)-    (Text, Text, Text) ->-    -- | delay before message send-    Int ->-    -- | protobuf message to send-    a ->-    m ()   sendEgressMsg ::     Message a =>     -- | egress address (namespace, type)@@ -163,22 +175,22 @@     -- | protobuf message to send (should be a Kafka or Kinesis protobuf record)     a ->     m ()--  sendByteMsg ::+  sendMsg ::     Serde a =>     -- | Function address (namespace, type, id)-    (Text, Text, Text) ->+    Address ->     -- | message to send     a ->     m ()-  sendByteMsgDelay ::+  sendMsgDelay ::     Serde a =>     -- | Function address (namespace, type, id)-    (Text, Text, Text) ->+    Address ->     -- | delay before message send-    Int ->+    NominalDiffTime ->     -- | message to send     a ->+    -- | returns cancelation token with which delivery of the message could be canceled     m ()  instance StatefulFunc s (Function s) where@@ -187,54 +199,29 @@   insideCtx func = func <$> getCtx   getCtx = gets functionStateCtx   setCtx new = modify (\old -> old {functionStateCtx = new, functionStateMutated = True})-  modifyCtx mutator = mutator <$> getCtx >>= setCtx-  sendMsg (namespace, funcType, id') msg = do-    invocations <- gets functionStateInvocations-    modify (\old -> old {functionStateInvocations = invocations Seq.:|> invocation})-    where-      target :: PR.Address-      target =-        defMessage-          & PR.namespace .~ namespace-          & PR.type' .~ funcType-          & PR.id .~ id'-      invocation :: PR.FromFunction'Invocation-      invocation =-        defMessage-          & PR.target .~ target-          & PR.argument .~ Any.pack msg-  sendMsgDelay (namespace, funcType, id') delay msg = do-    invocations <- gets functionStateDelayedInvocations-    modify (\old -> old {functionStateDelayedInvocations = invocations Seq.:|> invocation})-    where-      target :: PR.Address-      target =-        defMessage-          & PR.namespace .~ namespace-          & PR.type' .~ funcType-          & PR.id .~ id'-      invocation :: PR.FromFunction'DelayedInvocation-      invocation =-        defMessage-          & PR.delayInMs .~ fromIntegral delay-          & PR.target .~ target-          & PR.argument .~ Any.pack msg+  modifyCtx mutator = getCtx >>= setCtx . mutator+   sendEgressMsg (namespace, egressType) msg = do     egresses <- gets functionStateEgressMessages     modify (\old -> old {functionStateEgressMessages = egresses Seq.:|> egressMsg})     where+      wmsg = ProtoSerde msg       egressMsg :: PR.FromFunction'EgressMessage       egressMsg =         defMessage           & PR.egressNamespace .~ namespace           & PR.egressType .~ egressType-          & PR.argument .~ Any.pack msg-  sendByteMsg (namespace, funcType, id') msg = do+          & PR.argument .~ tpValue+      tpValue =+        defMessage+          & PR.typename .~ tpName (pure wmsg)+          & PR.hasValue .~ True          +          & PR.value .~ serializeBytes wmsg++  sendMsg (Address' namespace funcType id') msg = do     invocations <- gets functionStateInvocations     modify (\old -> old {functionStateInvocations = invocations Seq.:|> invocation})     where-      argument :: Any-      argument = defMessage & Any.value .~ serializeBytes msg       target :: PR.Address       target =         defMessage@@ -245,13 +232,17 @@       invocation =         defMessage           & PR.target .~ target-          & PR.argument .~ argument-  sendByteMsgDelay (namespace, funcType, id') delay msg = do+          & PR.argument .~ tpValue+      tpValue =+        defMessage+          & PR.typename .~ tpName (pure msg)+          & PR.hasValue .~ True+          & PR.value .~ serializeBytes msg++  sendMsgDelay (Address' namespace funcType id') delay msg = do     invocations <- gets functionStateDelayedInvocations     modify (\old -> old {functionStateDelayedInvocations = invocations Seq.:|> invocation})     where-      argument :: Any-      argument = defMessage & Any.value .~ serializeBytes msg       target :: PR.Address       target =         defMessage@@ -261,13 +252,28 @@       invocation :: PR.FromFunction'DelayedInvocation       invocation =         defMessage-          & PR.delayInMs .~ fromIntegral delay+          & PR.delayInMs .~ round (delay * 1000)           & PR.target .~ target-          & PR.argument .~ argument+          & PR.argument .~ tpValue+      tpValue =+        defMessage+          & PR.typename .~ tpName (pure msg)+          & PR.hasValue .~ True+          & PR.value .~ serializeBytes msg +-- | Convinience function to send protobuf messages+sendProtoMsg :: (StatefulFunc s m, Message a) => Address -> a -> m ()+sendProtoMsg addr = sendMsg addr . ProtoSerde++-- | Convinience function to send delayed protobuf messages+sendProtoMsgDelay :: (StatefulFunc s m, Message a) => Address -> NominalDiffTime -> a -> m ()+sendProtoMsgDelay addr delay = sendMsgDelay addr delay . ProtoSerde+ data FlinkError   = MissingInvocationBatch   | ProtodeserializeBytesError String+  | InvalidTypePassedError Text Text+  | EmptyArgumentPassed   | StateDecodeError String   | MessageDecodeError String   | ProtoMessageDecodeError UnpackError@@ -275,63 +281,53 @@   deriving (Show, Eq)  -- | Convenience function for wrapping state in newtype for JSON serialization-jsonState :: Json s => (a -> Function s ()) -> a -> Function (JsonSerde s) ()-jsonState f a = jsonWrapper-  where-    jsonWrapper = Function (ExceptT (StateT (\s -> ReaderT (\env -> (fmap . fmap) JsonSerde <$> runner (f a) env (getJson <$> s)))))-    runner res env state = runReaderT (runStateT (runExceptT $ runFunction res) state) env+jsonState :: Json s => Function s () -> Function (JsonSerde s) ()+jsonState = coerce  -- | Convenience function for wrapping state in newtype for Protobuf serialization-protoState :: Message s => (a -> Function s ()) -> a -> Function (ProtoSerde s) ()-protoState f a = protoWrapper-  where-    protoWrapper = Function (ExceptT (StateT (\s -> ReaderT (\env -> (fmap . fmap) ProtoSerde <$> runner (f a) env (getMessage <$> s)))))-    runner res env state = runReaderT (runStateT (runExceptT $ runFunction res) state) env+protoState :: Message s => Function s () -> Function (ProtoSerde s) ()+protoState = coerce +-- | Tries to unwrap typed value, possibly throwing FlinkError on broken input+unwrapA :: forall a m. (Serde a, MonadError FlinkError m) => PR.TypedValue -> m (Maybe a)+unwrapA arg = let+      atp = arg ^. PR.typename+      ctp = tpName @a Proxy+      in if not (arg^. PR.hasValue) then pure Nothing else+           if atp /= ctp+              then throwError (InvalidTypePassedError ctp atp)+              else Just <$> (liftEither . mapLeft MessageDecodeError $ deserializeBytes @a (arg ^. PR.value))+ -- | Takes a function taking an arbitrary state type and converts it to take 'ByteString's. -- This allows each function in the 'FunctionTable' to take its own individual type of state and just expose -- a function accepting 'ByteString' to the library code.-flinkWrapper :: Serde s => (Any -> Function s ()) -> ByteString -> Env -> PR.ToFunction'InvocationBatchRequest -> IO (Either FlinkError (FunctionState ByteString))-flinkWrapper func initialContext env invocationBatch = runExceptT $ do-  deserializeBytesdContext <- liftEither $ mapLeft StateDecodeError $ deserializeBytes initialContext-  (err, finalState) <- liftIO $ runner (newState deserializeBytesdContext)-  liftEither err-  return $ serializeBytes <$> finalState+flinkWrapper :: forall a s. (Serde a, Serde s) => s -> Expiration -> (a -> Function s ()) -> FuncExec+flinkWrapper s0 expr func env invocationBatch = runExceptT $ do+  (eiRes, _) <- liftIO $ runner (newState s0)+  liftEither eiRes   where-    runner state = runReaderT (runStateT (runExceptT $ runFunction runWithCtx) state) env+    passedArgs = invocationBatch ^.. PR.invocations . traversed . PR.argument+    mbInitCtx = (PR.state . traversed . filtered ((== "flink_state") . (^. PR.stateName)) . PR.stateValue) +      `firstOf` invocationBatch     runWithCtx = do-      defaultCtx <- gets functionStateCtx-      let initialCtx = getInitialCtx defaultCtx (listToMaybe $ invocationBatch ^. PR.state)-      case initialCtx of-        Left err -> throwError err-        Right ctx -> setInitialCtx ctx-      mapM_ func ((^. PR.argument) <$> invocationBatch ^. PR.invocations)--    getInitialCtx def pv = handleEmptyState def $ (^. PR.stateValue) <$> pv-    handleEmptyState def state' = case state' of-      Just "" -> return def-      Just other -> mapLeft StateDecodeError $ deserializeBytes other-      Nothing -> return def---- | Deserializes input messages as arbitrary bytes by extracting them out of the protobuf Any--- and ignoring the type since that's protobuf specific-serdeInput :: (Serde s, Serde a, StatefulFunc s m, MonadError FlinkError m, MonadReader Env m) => (a -> m b) -> Any -> m b-serdeInput f input = f =<< liftEither (mapLeft MessageDecodeError $ deserializeBytes (input ^. Any.value))---- | Deserializes input messages as arbitrary bytes by extracting them out of the protobuf Any--- and ignoring the type since that's protobuf specific-jsonInput :: (Serde s, Json a, StatefulFunc s m, MonadError FlinkError m, MonadReader Env m) => (a -> m b) -> Any -> m b-jsonInput f = serdeInput wrapJson-  where-    wrapJson (JsonSerde msg) = f msg---- | Deserializes input messages by unpacking the protobuf Any into the expected type.--- If you are passing messages via protobuf, this is much more typesafe than 'serdeInput'-protoInput :: (Serde s, Message a, StatefulFunc s m, MonadError FlinkError m, MonadReader Env m) => (a -> m b) -> Any -> m b-protoInput f input = f =<< liftEither (mapLeft ProtoMessageDecodeError $ Any.unpack input)+      case mbInitCtx of +        Nothing -> return $ IncompleteContext expr (tpName @s Proxy) -- if state was not propagated to the function - shorcut to incomplete context reponse+        Just tv -> do +          mbCtx <- unwrapA @s tv+          case mbCtx of+            Nothing -> pure () -- if null state value was propagated+            Just s1 -> setInitialCtx s1+          mbArgs <- traverse (unwrapA @a) passedArgs+          args <- traverse (maybe (throwError EmptyArgumentPassed) pure) mbArgs+          mapM_ func args+          gets (UpdatedState . fmap outS)+    runner state = runReaderT (runStateT (runExceptT $ runFunction runWithCtx) state) env+    outS fstate = defMessage+        & PR.typename .~ tpName (pure fstate)+        & PR.value .~ serializeBytes fstate -createFlinkResp :: FunctionState ByteString -> FromFunction-createFlinkResp (FunctionState state mutated invocations delayedInvocations egresses) =+createFlinkResp :: FuncRes -> FromFunction+createFlinkResp (UpdatedState (FunctionState state mutated invocations delayedInvocations egresses)) =   defMessage & PR.invocationResult     .~ ( defMessage            & PR.stateMutations .~ toList stateMutations@@ -348,6 +344,20 @@           & PR.stateValue .~ state         | mutated       ]+createFlinkResp (IncompleteContext (Expiration mode expireTime) typeName) = +  defMessage & PR.incompleteInvocationContext .~ ( +    defMessage & PR.missingValues .~ [+      defMessage+        & PR.stateName .~ "flink_state"+        & PR.typeTypename .~ typeName+        & PR.expirationSpec .~ (+          defMessage +            & PR.expireAfterMillis .~ round (expireTime * 1000.0)+            & PR.mode .~ pbmode mode)+      ])+  where pbmode NONE = PR.FromFunction'ExpirationSpec'NONE +        pbmode AFTER_CALL = PR.FromFunction'ExpirationSpec'AFTER_INVOKE  +        pbmode AFTER_WRITE = PR.FromFunction'ExpirationSpec'AFTER_WRITE    type FlinkApi =   "statefun" :> ReqBody '[Proto] ToFunction :> Post '[Proto] FromFunction@@ -363,14 +373,14 @@ flinkServer :: FunctionTable -> Server FlinkApi flinkServer functions toFunction = do   batch <- getBatch toFunction-  ((initialCtx, function), (namespace, type', id')) <- findFunc (batch ^. PR.target)-  result <- liftIO $ function initialCtx (Env namespace type' id') batch+  (function, (namespace, type', id')) <- findFunc (batch ^. PR.target)+  result <- liftIO $ function (Env namespace type' id') batch   finalState <- liftEither $ mapLeft flinkErrToServant result   return $ createFlinkResp finalState   where     getBatch input = maybe (throwError $ flinkErrToServant MissingInvocationBatch) return (input ^? PR.maybe'request . _Just . PR._ToFunction'Invocation')     findFunc addr = do-      res <- maybe (throwError $ flinkErrToServant $ NoSuchFunction (namespace, type')) return (Map.lookup (namespace, type') functions)+      res <- maybe (throwError $ flinkErrToServant $ NoSuchFunction (namespace, type')) return (Map.lookup (FuncType namespace type') functions)       return (res, address)       where         address@(namespace, type', _) = (addr ^. PR.namespace, addr ^. PR.type', addr ^. PR.id)@@ -383,3 +393,5 @@   MessageDecodeError msg -> err400 {errBody = "Failed to decode message " <> BSL.pack msg}   ProtoMessageDecodeError msg -> err400 {errBody = "Failed to decode message " <> BSL.pack (show msg)}   NoSuchFunction (namespace, type') -> err400 {errBody = "No such function " <> T.encodeUtf8 (fromStrict namespace) <> T.encodeUtf8 (fromStrict type')}+  InvalidTypePassedError expected passed -> err400 {errBody = "Expected type " <> T.encodeUtf8 (fromStrict expected) <> ", got " <>  T.encodeUtf8 (fromStrict passed)}+  EmptyArgumentPassed -> err400 {errBody = "Empty argument was passed to the function" }

@@ -17,23 +17,24 @@         modifyCtx,         sendMsg,         sendMsgDelay,-        sendEgressMsg,-        sendByteMsg,-        sendByteMsgDelay+        sendEgressMsg       ),     flinkWrapper,     createApp,     flinkServer,     flinkApi,+    Address(..),+    FuncType (..),     Serde (..),     FunctionTable,     JsonSerde (..),     ProtoSerde (..),+    Expiration(..),+    ExpirationMode(..),     jsonState,     protoState,-    serdeInput,-    protoInput,-    jsonInput,+    sendProtoMsg,+    sendProtoMsgDelay   ) where @@ -43,12 +44,16 @@       ProtoSerde(..),       Serde(..),       FunctionTable,+      Expiration(..),+      ExpirationMode(..),+      Address(..),+      FuncType (..),       jsonState,       protoState,       flinkApi,       createApp,       flinkServer,-      protoInput,-      serdeInput,       flinkWrapper,-      jsonInput )+      sendProtoMsg,+      sendProtoMsgDelay+ )