// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

// SPDX-License-Identifier: Apache-2.0

package rapidcore

import (

"bytes"

"context"

"errors"

"fmt"

"io"

"io/ioutil"

"math"

"net/http"

"sync"

"time"

"go.amzn.com/lambda/core/directinvoke"

"go.amzn.com/lambda/core/statejson"

"go.amzn.com/lambda/interop"

"go.amzn.com/lambda/metering"

"github.com/google/uuid"

log "github.com/sirupsen/logrus"

)

const (

autoresetReasonTimeout = "Timeout"

autoresetReasonReserveFail = "ReserveFail"

autoresetReasonReleaseFail = "ReleaseFail"

standaloneVersionID = "1"

resetDefaultTimeoutMs = 2000

)

type rapidPhase int

const (

phaseIdle rapidPhase = iota

phaseInitializing

phaseInvoking

)

type runtimeState int

const (

runtimeNotStarted = iota

runtimeInitStarted

runtimeInitError

runtimeInitComplete

runtimeInitFailed

runtimeInvokeResponseSent

runtimeInvokeError

runtimeReady

runtimeInvokeComplete

)

type DoneWithState struct {

*interop.Done

State statejson.InternalStateDescription

}

func (s *DoneWithState) String() string {

return fmt.Sprintf("%v %v", *s.Done, string(s.State.AsJSON()))

}

type InvokeContext struct {

Token interop.Token

ReplySent bool

ReplyStream http.ResponseWriter

Direct bool

}

type Server struct {

InternalStateGetter interop.InternalStateGetter

invokeChanOut chan *interop.Invoke

startChanOut chan *interop.Start

resetChanOut chan *interop.Reset

shutdownChanOut chan *interop.Shutdown

errorChanOut chan error

sendRunningChan chan *interop.Running

sendResponseChan chan struct{}

doneChan chan *interop.Done

InitDoneChan chan DoneWithState

InvokeDoneChan chan DoneWithState

ResetDoneChan chan *interop.Done

ShutdownDoneChan chan *interop.Done

mutex sync.Mutex

invokeCtx *InvokeContext

invokeTimeout time.Duration

reservationContext context.Context

reservationCancel func()

rapidPhase rapidPhase

runtimeState runtimeState

}

func (s *Server) StartAcceptingDirectInvokes() error {

return nil

}

func (s *Server) setRapidPhase(phase rapidPhase) {

s.mutex.Lock()

defer s.mutex.Unlock()

s.rapidPhase = phase

}

func (s *Server) getRapidPhase() rapidPhase {

s.mutex.Lock()

defer s.mutex.Unlock()

return s.rapidPhase

}

func (s *Server) setRuntimeState(state runtimeState) {

s.mutex.Lock()

defer s.mutex.Unlock()

s.runtimeState = state

}

func (s *Server) getRuntimeState() runtimeState {

s.mutex.Lock()

defer s.mutex.Unlock()

return s.runtimeState

}

func (s *Server) SetInvokeTimeout(timeout time.Duration) {

s.mutex.Lock()

defer s.mutex.Unlock()

s.invokeTimeout = timeout

}

func (s *Server) GetInvokeTimeout() time.Duration {

s.mutex.Lock()

defer s.mutex.Unlock()

return s.invokeTimeout

}

func (s *Server) GetInvokeContext() *InvokeContext {

s.mutex.Lock()

defer s.mutex.Unlock()

ctx := *s.invokeCtx

return &ctx

}

func (s *Server) setNewInvokeContext(invokeID string, traceID, lambdaSegmentID string) (*ReserveResponse, error) {

s.mutex.Lock()

defer s.mutex.Unlock()

if s.invokeCtx != nil {

return nil, ErrAlreadyReserved

}

s.invokeCtx = &InvokeContext{

Token: interop.Token{

ReservationToken: uuid.New().String(),

InvokeID: invokeID,

VersionID: standaloneVersionID,

FunctionTimeout: s.invokeTimeout,

TraceID: traceID,

LambdaSegmentID: lambdaSegmentID,

InvackDeadlineNs: math.MaxInt64, // no INVACK in standalone

},

}

resp := &ReserveResponse{

Token: s.invokeCtx.Token,

}

s.reservationContext, s.reservationCancel = context.WithCancel(context.Background())

return resp, nil

}

// Reserve allocates invoke context

func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveResponse, error) {

invokeID := uuid.New().String()

if len(id) > 0 {

invokeID = id

}

resp, err := s.setNewInvokeContext(invokeID, traceID, lambdaSegmentID)

if err != nil {

return nil, err

}

resp.InternalState, err = s.waitInit()

return resp, err

}

func (s *Server) waitInit() (*statejson.InternalStateDescription, error) {

for {

select {

case doneWithState, chanOpen := <-s.InitDoneChan:

if !chanOpen {

// init only happens once

return nil, ErrInitAlreadyDone

}

close(s.InitDoneChan) // this was first call to reserve

if s.getRuntimeState() == runtimeInitFailed {

return &doneWithState.State, ErrInitError

}

if len(doneWithState.ErrorType) > 0 {

log.Errorf("INIT DONE failed: %s", doneWithState.ErrorType)

return &doneWithState.State, ErrInitDoneFailed

}

return &doneWithState.State, nil

case <-s.reservationContext.Done():

return nil, ErrReserveReservationDone

}

}

}

func (s *Server) setReplyStream(w http.ResponseWriter, direct bool) (string, error) {

s.mutex.Lock()

defer s.mutex.Unlock()

if s.invokeCtx == nil {

return "", ErrNotReserved

}

if s.invokeCtx.ReplySent {

return "", ErrAlreadyReplied

}

if s.invokeCtx.ReplyStream != nil {

return "", ErrAlreadyInvocating

}

s.invokeCtx.ReplyStream = w

s.invokeCtx.Direct = direct

return s.invokeCtx.Token.InvokeID, nil

}

// Release closes the invocation, making server ready for reserve again

func (s *Server) Release() error {

s.mutex.Lock()

defer s.mutex.Unlock()

if s.invokeCtx == nil {

return ErrNotReserved

}

if s.reservationCancel != nil {

s.reservationCancel()

}

s.invokeCtx = nil

return nil

}

// GetCurrentInvokeID

func (s *Server) GetCurrentInvokeID() string {

s.mutex.Lock()

defer s.mutex.Unlock()

if s.invokeCtx == nil {

return ""

}

return s.invokeCtx.Token.InvokeID

}

// SetInternalStateGetter is used to set callback which returnes internal state for /test/internalState request

func (s *Server) SetInternalStateGetter(cb interop.InternalStateGetter) {

s.InternalStateGetter = cb

}

// StartChan returns Start emitter

func (s *Server) StartChan() <-chan *interop.Start {

return s.startChanOut

}

// InvokeChan returns Invoke emitter

func (s *Server) InvokeChan() <-chan *interop.Invoke {

return s.invokeChanOut

}

// ResetChan returns Reset emitter

func (s *Server) ResetChan() <-chan *interop.Reset {

return s.resetChanOut

}

// ShutdownChan returns Shutdown emitter

func (s *Server) ShutdownChan() <-chan *interop.Shutdown {

return s.shutdownChanOut

}

// InvalidMessageChan emits errors if there was something we could not parse

func (s *Server) TransportErrorChan() <-chan error {

return s.errorChanOut

}

func (s *Server) sendResponseUnsafe(invokeID string, contentType string, status int, payload io.Reader) error {

if s.invokeCtx == nil || invokeID != s.invokeCtx.Token.InvokeID {

return interop.ErrInvalidInvokeID

}

if s.invokeCtx.ReplySent {

return interop.ErrResponseSent

}

if s.invokeCtx.ReplyStream == nil {

return fmt.Errorf("ReplyStream not available")

}

// TODO: earlier, status was set to 500 if runtime called /invocation/error. However, the integration

// tests do not differentiate between /invocation/error and /invocation/response, but they check the error type:

// To identify user-errors, we should also allowlist custom errortypes and propagate them via headers.

// s.invokeCtx.ReplyStream.WriteHeader(status)

var reportedErr error

if s.invokeCtx.Direct {

if err := directinvoke.SendDirectInvokeResponse(map[string]string{"Content-Type": contentType}, payload, s.invokeCtx.ReplyStream); err != nil {

// TODO: Do we need to drain the reader in case of a large payload and connection reuse?

log.Errorf("Failed to write response to %s: %s", invokeID, err)

flusher, ok := s.invokeCtx.ReplyStream.(http.Flusher)

if !ok {

log.Error("Failed to flush response")

}

flusher.Flush()

reportedErr = err

}

} else {

data, err := ioutil.ReadAll(payload)

if err != nil {

return fmt.Errorf("Failed to read response on %s: %s", invokeID, err)

}

if len(data) > interop.MaxPayloadSize {

return &interop.ErrorResponseTooLarge{

ResponseSize: len(data),

MaxResponseSize: interop.MaxPayloadSize,

}

}

s.invokeCtx.ReplyStream.Header().Add("Content-Type", contentType)

if _, err := s.invokeCtx.ReplyStream.Write(data); err != nil {

return fmt.Errorf("Failed to write response to %s: %s", invokeID, err)

}

}

s.sendResponseChan <- struct{}{}

s.invokeCtx.ReplySent = true

s.invokeCtx.Direct = false

return reportedErr

}

func (s *Server) SendResponse(invokeID string, contentType string, reader io.Reader) error {

s.setRuntimeState(runtimeInvokeResponseSent)

s.mutex.Lock()

defer s.mutex.Unlock()

return s.sendResponseUnsafe(invokeID, contentType, http.StatusOK, reader)

}

func (s *Server) CommitResponse() error { return nil }

func (s *Server) SendRunning(run *interop.Running) error {

s.setRuntimeState(runtimeInitStarted)

s.sendRunningChan <- run

return nil

}

func (s *Server) SendErrorResponse(invokeID string, resp *interop.ErrorResponse) error {

switch s.getRapidPhase() {

case phaseInitializing:

s.setRuntimeState(runtimeInitError)

return nil

case phaseInvoking:

// This branch can also occur during a suppressed init error, which is reported as invoke error

s.setRuntimeState(runtimeInvokeError)

s.mutex.Lock()

defer s.mutex.Unlock()

return s.sendResponseUnsafe(invokeID, resp.ContentType, http.StatusInternalServerError, bytes.NewReader(resp.Payload))

default:

panic("received unexpected error response outside invoke or init phases")

}

}

func (s *Server) SendDone(done *interop.Done) error {

s.doneChan <- done

return nil

}

func (s *Server) SendDoneFail(doneFail *interop.DoneFail) error {

s.doneChan <- &interop.Done{

ErrorType: doneFail.ErrorType,

CorrelationID: doneFail.CorrelationID, // filipovi: correlationID is required to dispatch message into correct channel

Meta: doneFail.Meta,

}

return nil

}

func (s *Server) Reset(reason string, timeoutMs int64) (*statejson.ResetDescription, error) {

// pass reset to rapid

s.resetChanOut <- &interop.Reset{

Reason: reason,

DeadlineNs: deadlineNsFromTimeoutMs(timeoutMs),

CorrelationID: "resetCorrelationID",

}

// TODO do not block on reset, instead consume ResetDoneChan in waitForRelease handler,

// this will get us more aligned on async reset notification handling.

done := <-s.ResetDoneChan

s.Release()

if done.ErrorType != "" {

return nil, errors.New(string(done.ErrorType))

}

return &statejson.ResetDescription{ExtensionsResetMs: done.Meta.ExtensionsResetMs}, nil

}

func NewServer(ctx context.Context) *Server {

s := &Server{

startChanOut: make(chan *interop.Start),

invokeChanOut: make(chan *interop.Invoke),

errorChanOut: make(chan error),

resetChanOut: make(chan *interop.Reset),

shutdownChanOut: make(chan *interop.Shutdown),

sendRunningChan: make(chan *interop.Running),

sendResponseChan: make(chan struct{}),

doneChan: make(chan *interop.Done),

// These two channels are buffered, because they are depleted asynchronously (by reserve and waitUntilRelease) and we don't want to block in SendDone until they are called

InitDoneChan: make(chan DoneWithState, 1),

InvokeDoneChan: make(chan DoneWithState, 1),

ResetDoneChan: make(chan *interop.Done),

ShutdownDoneChan: make(chan *interop.Done),

}

go s.dispatchDone()

return s

}

func (s *Server) setInitDoneRuntimeState(done *interop.Done) {

if len(done.ErrorType) > 0 {

s.setRuntimeState(runtimeInitFailed) // donefail

} else {

s.setRuntimeState(runtimeInitComplete) // done

}

}

// Note, the dispatch loop below has potential to block, when

// channel is not drained. E.g. if test assumes sandbox init

// completion before dispatching reset, then reset will block

// until init channel is drained.

func (s *Server) dispatchDone() {

for {

done := <-s.doneChan

log.Debug("Dispatching DONE:", done.CorrelationID)

internalState := s.InternalStateGetter()

s.setRapidPhase(phaseIdle)

if done.CorrelationID == "initCorrelationID" {

s.setInitDoneRuntimeState(done)

s.InitDoneChan <- DoneWithState{Done: done, State: internalState}

} else if done.CorrelationID == "invokeCorrelationID" {

s.setRuntimeState(runtimeInvokeComplete)

s.InvokeDoneChan <- DoneWithState{Done: done, State: internalState}

} else if done.CorrelationID == "resetCorrelationID" {

s.setRuntimeState(runtimeNotStarted)

s.ResetDoneChan <- done

} else if done.CorrelationID == "shutdownCorrelationID" {

s.setRuntimeState(runtimeNotStarted)

s.ShutdownDoneChan <- done

} else {

panic("Received DONE without correlation ID")

}

}

}

func drainChannel(c chan DoneWithState) {

for {

select {

case dws := <-c:

log.Warnf("Discard DONE response: %s", dws.String())

break

default:

return

}

}

}

func (s *Server) Clear() {

// we do not drain InitDoneChannel, because Init is only done once during rapid lifetime

drainChannel(s.InvokeDoneChan)

s.Release()

}

func (s *Server) IsResponseSent() bool {

panic("unexpected call to unimplemented method in rapidcore: IsResponseSent()")

}

func (s *Server) SendRuntimeReady() error {

// only called when extensions are enabled

s.setRuntimeState(runtimeReady)

return nil

}

func deadlineNsFromTimeoutMs(timeoutMs int64) int64 {

mono := metering.Monotime()

return mono + timeoutMs*1000*1000

}

func (s *Server) Init(i *interop.Start, invokeTimeoutMs int64) {

s.SetInvokeTimeout(time.Duration(invokeTimeoutMs) * time.Millisecond)

s.startChanOut <- i

s.setRapidPhase(phaseInitializing)

<-s.sendRunningChan

log.Debug("Received RUNNING")

}

func (s *Server) FastInvoke(w http.ResponseWriter, i *interop.Invoke, direct bool) error {

invokeID, err := s.setReplyStream(w, direct)

if err != nil {

return err

}

s.setRapidPhase(phaseInvoking)

i.ID = invokeID

select {

case s.invokeChanOut <- i:

break

case <-s.sendResponseChan:

// we didn't pass invoke to rapid yet, but rapid already has written some response

// It can happend if runtime/agent crashed even before we passed invoke to it

return ErrInvokeResponseAlreadyWritten

}

select {

case <-s.sendResponseChan:

break

case <-s.reservationContext.Done():

return ErrInvokeReservationDone

}

return nil

}

func (s *Server) CurrentToken() *interop.Token {

s.mutex.Lock()

defer s.mutex.Unlock()

if s.invokeCtx == nil {

return nil

}

tok := s.invokeCtx.Token

return &tok

}

// Invoke is used by the Runtime Interface Emulator (Rapid Local)

// https://github.com/aws/aws-lambda-runtime-interface-emulator

func (s *Server) Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error {

resetCtx, resetCancel := context.WithCancel(context.Background())

defer resetCancel()

timeoutChan := make(chan error)

go func() {

select {

case <-time.After(s.GetInvokeTimeout()):

timeoutChan <- ErrInvokeTimeout

s.Reset(autoresetReasonTimeout, resetDefaultTimeoutMs)

case <-resetCtx.Done():

log.Debugf("execute finished, autoreset cancelled")

}

}()

reserveResp, err := s.Reserve(invoke.ID, "", "")

if err != nil {

switch err {

case ErrInitError:

// Simulate 'Suppressed Init' scenario

s.Reset(autoresetReasonReserveFail, resetDefaultTimeoutMs)

reserveResp, err = s.Reserve("", "", "")

if err == ErrInitAlreadyDone {

break

}

return err

case ErrInitDoneFailed, ErrTerminated:

s.Reset(autoresetReasonReserveFail, resetDefaultTimeoutMs)

return err

case ErrInitAlreadyDone:

// This is a valid response (e.g. for 2nd and 3rd invokes)

// TODO: switch on ReserveResponse status instead of err,

// since these are valid values

if s.InternalStateGetter == nil {

responseWriter.Write([]byte("error: internal state callback not set"))

return ErrInternalServerError

}

default:

return err

}

}

invoke.DeadlineNs = fmt.Sprintf("%d", metering.Monotime()+reserveResp.Token.FunctionTimeout.Nanoseconds())

invokeChan := make(chan error)

go func() {

if err := s.FastInvoke(responseWriter, invoke, false); err != nil {

invokeChan <- err

}

}()

releaseChan := make(chan error)

go func() {

_, err := s.AwaitRelease()

releaseChan <- err

}()

// TODO: verify the order of channel receives. When timeouts happen, Reset()

// is called first, which also does Release() => this may signal a type

// Err<*>ReservationDone error to the non-timeout channels. This is currently

// handled by the http handler, which returns GatewayTimeout for reservation errors

// too. However, Timeouts should ideally be only represented by ErrInvokeTimeout.

select {

case err = <-invokeChan:

case err = <-timeoutChan:

case err = <-releaseChan:

if err != nil {

s.Reset(autoresetReasonReleaseFail, resetDefaultTimeoutMs)

}

}

return err

}

func (s *Server) AwaitRelease() (*statejson.InternalStateDescription, error) {

select {

case doneWithState := <-s.InvokeDoneChan:

if len(doneWithState.ErrorType) > 0 {

log.Errorf("Invoke DONE failed: %s", doneWithState.ErrorType)

return nil, ErrInvokeDoneFailed

}

s.Release()

return &doneWithState.State, nil

case <-s.reservationContext.Done():

return nil, ErrReleaseReservationDone

}

}

func (s *Server) Shutdown(shutdown *interop.Shutdown) *statejson.InternalStateDescription {

s.shutdownChanOut <- shutdown

<-s.ShutdownDoneChan

state := s.InternalStateGetter()

return &state

}

func (s *Server) InternalState() (*statejson.InternalStateDescription, error) {

if s.InternalStateGetter == nil {

return nil, errors.New("InternalStateGetterNotSet")

}

state := s.InternalStateGetter()

return &state, nil

}