@@ -11,6 +11,11 @@ import (
1111 "google.golang.org/protobuf/proto"
1212)
1313
14+ type invocation interface {
15+ awaitOutput (timeout * time.Duration ) (any , error )
16+ retry (retryCount uint32 ) error
17+ }
18+
1419// controlPlaneInvocation implements the invocation interface.
1520type controlPlaneInvocation struct {
1621 FunctionCallId string
@@ -52,7 +57,7 @@ func controlPlaneInvocationFromFunctionCallId(ctx context.Context, functionCallI
5257}
5358
5459func (c * controlPlaneInvocation ) awaitOutput (timeout * time.Duration ) (any , error ) {
55- return pollFunctionOutput (c .ctx , c .FunctionCallId , timeout )
60+ return pollFunctionOutput (c .ctx , c .getOutput , timeout )
5661}
5762
5863func (c * controlPlaneInvocation ) retry (retryCount uint32 ) error {
@@ -75,8 +80,102 @@ func (c *controlPlaneInvocation) retry(retryCount uint32) error {
7580 return nil
7681}
7782
78- // Poll for outputs for a given FunctionCall ID.
79- func pollFunctionOutput (ctx context.Context , functionCallId string , timeout * time.Duration ) (any , error ) {
83+ // getOutput fetches the output for the current function call with a timeout in milliseconds.
84+ func (c * controlPlaneInvocation ) getOutput (timeout time.Duration ) (* pb.FunctionGetOutputsItem , error ) {
85+ response , err := client .FunctionGetOutputs (c .ctx , pb.FunctionGetOutputsRequest_builder {
86+ FunctionCallId : c .FunctionCallId ,
87+ MaxValues : 1 ,
88+ Timeout : float32 (timeout .Seconds ()),
89+ LastEntryId : "0-0" ,
90+ ClearOnSuccess : true ,
91+ RequestedAt : timeNowSeconds (),
92+ }.Build ())
93+ if err != nil {
94+ return nil , fmt .Errorf ("FunctionGetOutputs failed: %w" , err )
95+ }
96+ outputs := response .GetOutputs ()
97+ if len (outputs ) > 0 {
98+ return outputs [0 ], nil
99+ }
100+ return nil , nil
101+ }
102+
103+ // InputPlaneInvocation implements the Invocation interface for the input plane.
104+ type inputPlaneInvocation struct {
105+ client pb.ModalClientClient
106+ functionId string
107+ input * pb.FunctionPutInputsItem
108+ attemptToken string
109+ ctx context.Context
110+ }
111+
112+ // CreateInputPlaneInvocation creates a new InputPlaneInvocation by starting an attempt.
113+ func createInputPlaneInvocation (ctx context.Context , inputPlaneURL string , functionId string , input * pb.FunctionInput ) (* inputPlaneInvocation , error ) {
114+ functionPutInputsItem := pb.FunctionPutInputsItem_builder {
115+ Idx : 0 ,
116+ Input : input ,
117+ }.Build ()
118+ client , err := getOrCreateInputPlaneClient (inputPlaneURL )
119+ if err != nil {
120+ return nil , err
121+ }
122+ attemptStartResp , err := client .AttemptStart (ctx , pb.AttemptStartRequest_builder {
123+ FunctionId : functionId ,
124+ Input : functionPutInputsItem ,
125+ }.Build ())
126+ if err != nil {
127+ return nil , err
128+ }
129+ return & inputPlaneInvocation {
130+ client : client ,
131+ functionId : functionId ,
132+ input : functionPutInputsItem ,
133+ attemptToken : attemptStartResp .GetAttemptToken (),
134+ ctx : ctx ,
135+ }, nil
136+ }
137+
138+ // awaitOutput waits for the output with an optional timeout.
139+ func (i * inputPlaneInvocation ) awaitOutput (timeout * time.Duration ) (any , error ) {
140+ return pollFunctionOutput (i .ctx , i .getOutput , timeout )
141+ }
142+
143+ // getOutput fetches the output for the current attempt.
144+ func (i * inputPlaneInvocation ) getOutput (timeout time.Duration ) (* pb.FunctionGetOutputsItem , error ) {
145+ resp , err := i .client .AttemptAwait (i .ctx , pb.AttemptAwaitRequest_builder {
146+ AttemptToken : i .attemptToken ,
147+ RequestedAt : timeNowSeconds (),
148+ TimeoutSecs : float32 (timeout .Seconds ()),
149+ }.Build ())
150+ if err != nil {
151+ return nil , fmt .Errorf ("AttemptAwait failed: %w" , err )
152+ }
153+ return resp .GetOutput (), nil
154+ }
155+
156+ // retry retries the invocation.
157+ func (i * inputPlaneInvocation ) retry (retryCount uint32 ) error {
158+ // We ignore retryCount - it is used only by controlPlaneInvocation.
159+ resp , err := i .client .AttemptRetry (context .Background (), pb.AttemptRetryRequest_builder {
160+ FunctionId : i .functionId ,
161+ Input : i .input ,
162+ AttemptToken : i .attemptToken ,
163+ }.Build ())
164+ if err != nil {
165+ return err
166+ }
167+ i .attemptToken = resp .GetAttemptToken ()
168+ return nil
169+ }
170+
171+ // getOutput is a function type that takes a timeout and returns a FunctionGetOutputsItem or nil, and an error.
172+ // Used by `pollForOutputs` to fetch from either the control plane or the input plane, depending on the implementation.
173+ type getOutput func (timeout time.Duration ) (* pb.FunctionGetOutputsItem , error )
174+
175+ // pollFunctionOutput repeatedly tries to fetch an output using the provided `getOutput` function, and the specified
176+ // timeout value. We use a timeout value of 55 seconds if the caller does not specify a timeout value, or if the
177+ // specified timeout value is greater than 55 seconds.
178+ func pollFunctionOutput (ctx context.Context , getOutput getOutput , timeout * time.Duration ) (any , error ) {
80179 startTime := time .Now ()
81180 pollTimeout := outputsTimeout
82181 if timeout != nil {
@@ -85,23 +184,14 @@ func pollFunctionOutput(ctx context.Context, functionCallId string, timeout *tim
85184 }
86185
87186 for {
88- response , err := client .FunctionGetOutputs (ctx , pb.FunctionGetOutputsRequest_builder {
89- FunctionCallId : functionCallId ,
90- MaxValues : 1 ,
91- Timeout : float32 (pollTimeout .Seconds ()),
92- LastEntryId : "0-0" ,
93- ClearOnSuccess : true ,
94- RequestedAt : timeNowSeconds (),
95- }.Build ())
187+ output , err := getOutput (pollTimeout )
96188 if err != nil {
97- return nil , fmt . Errorf ( "FunctionGetOutputs failed: %w" , err )
189+ return nil , err
98190 }
99-
100191 // Output serialization may fail if any of the output items can't be deserialized
101192 // into a supported Go type. Users are expected to serialize outputs correctly.
102- outputs := response .GetOutputs ()
103- if len (outputs ) > 0 {
104- return processResult (ctx , outputs [0 ].GetResult (), outputs [0 ].GetDataFormat ())
193+ if output != nil {
194+ return processResult (ctx , output .GetResult (), output .GetDataFormat ())
105195 }
106196
107197 if timeout != nil {
0 commit comments