examples/step-function/main.go

// Command step-function is a Lambda used as a Step Functions task. Each
// invocation receives a state blob, mutates it, and returns the new state.
// lambdalog logs every transition with the execution id so runs can be
// reconstructed in CloudWatch Logs Insights.
//
// See mercemay.top/src/lambdalog/examples/step-function/.
package main

import (
	"context"
	"encoding/json"
	"errors"

	awslambda "github.com/aws/aws-lambda-go/lambda"

	adapter "mercemay.top/src/lambdalog/adapters/lambda"
)

type state struct {
	ExecutionID string `json:"execution_id"`
	Step        string `json:"step"`
	Attempts    int    `json:"attempts"`
	Done        bool   `json:"done"`
}

type logger struct{}

func (logger) Info(string, ...adapter.Field)        {}
func (logger) Error(string, ...adapter.Field)       {}
func (logger) With(...adapter.Field) adapter.Logger { return logger{} }

var errExhausted = errors.New("step-function: too many attempts")

func transition(_ context.Context, s state) (state, error) {
	s.Attempts++
	switch s.Step {
	case "", "start":
		s.Step = "processing"
	case "processing":
		s.Step = "finalising"
	case "finalising":
		s.Step = "done"
		s.Done = true
	default:
		return s, errExhausted
	}
	if s.Attempts > 10 {
		return s, errExhausted
	}
	return s, nil
}

func main() {
	h := adapter.HandlerFunc(logger{}, transition)
	awslambda.Start(h)
}

// simulate runs transition until Done or errExhausted. Tests use it to
// validate the state machine without deploying to AWS.
func simulate(initial state) (state, error) {
	cur := initial
	for !cur.Done {
		next, err := transition(context.Background(), cur)
		if err != nil {
			return next, err
		}
		// prevent infinite loops if transition returns the same state
		if next.Step == cur.Step && !next.Done {
			return next, errExhausted
		}
		cur = next
	}
	b, _ := json.Marshal(cur)
	_ = b
	return cur, nil
}