Repository: atlassian/gostatsd
Pull request #421 (2021-11-19 01:12:14)

Adding GostatsD Lambda Extension
Starting this off as a draft since it is _mostly_ complete. 

There is a few things I want to double check such as:
- [x] Forwarder will flush on context.Done()
- [x] release assets to github
- [x] The use if a virtual clock
- [ ] updated docs on how to use the extension
- [ ] a local integration test to validate work

However, since the logic is done. There is nothing stopping the current work from being reviewed.

MovieStoreGuy: Moving Releases to own Pull Request since that this draft is big
MovieStoreGuy: Forwarder to flush all metrics within a time interval is done here https://github.com/atlassian/gostatsd/pull/422/files

Patch:
diff --git a/.gitignore b/.gitignore
index b0437cd..05c58c2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 /build/bin/
+/build/release/
 /vendor/
 /tools/
 /coverage.out
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ccbb249..da3b282 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,7 +15,7 @@
 36.0.0
 ------
 - Bump go version to 1.20.6
-- Updated internal metrics to use `Report` and change on deployment mode. 
+- Updated internal metrics to use `Report` and change on deployment mode.
 - Updated METRICS.md for impacted values
 
 35.2.2
diff --git a/Makefile b/Makefile
index 516b348..10ffeb1 100644
--- a/Makefile
+++ b/Makefile
@@ -13,6 +13,7 @@ GOVERSION           := 1.21.4  # Go version needs to be the same in: CI config,
 GP                  := /gopath
 MAIN_PKG            := github.com/atlassian/gostatsd/cmd/gostatsd
 CLUSTER_PKG         := github.com/atlassian/gostatsd/cmd/cluster
+LAMBDA_PKG          := github.com/atlassian/gostatsd/cmd/lambda-extension
 PROTOBUF_VERSION    ?= 21.5
 PROJECT_ROOT_DIR    := $(realpath $(dir $(abspath $(lastword $(MAKEFILE_LIST)))))
 TOOLS_DIR           := $(PROJECT_ROOT_DIR)/.tools
@@ -52,13 +53,16 @@ build: pb/gostatsd.pb.go
 			$(PKG)
 
 build-gostatsd:
-	@$(MAKE) PKG=$(MAIN_PKG) 
+	@$(MAKE) build PKG=$(MAIN_PKG)
+
+build-lambda:
+	@$(MAKE) build PKG=$(LAMBDA_PKG) BINARY_NAME="lambda-extension"
 
 build-gostatsd-race:
 	@$(MAKE) build-gostatsd GOBUILD_OPTIONAL_FLAGS=-race
 
 build-cluster:
-	@$(MAKE) PKG=$(CLUSTER_PKG) BINARY_NAME="cluster"
+	@$(MAKE) build PKG=$(CLUSTER_PKG) BINARY_NAME="cluster"
 
 
 build-all: pb/gostatsd.pb.go tools
@@ -120,7 +124,7 @@ check-all: pb/gostatsd.pb.go
 	go install ./cmd/gostatsd
 	go install ./cmd/tester
 
-fuzz: $(TOOLS_DIR)/go-fuzz-build $(TOOLS_DIR)/go-fuzz 
+fuzz: $(TOOLS_DIR)/go-fuzz-build $(TOOLS_DIR)/go-fuzz
 	$(TOOLS_DIR)/go-fuzz-build github.com/atlassian/gostatsd/pkg/statsd
 	$(TOOLS_DIR)/go-fuzz -bin=./statsd-fuzz.zip -workdir=test_fixtures/lexer_fuzz
 
diff --git a/cmd/lambda-extension/main.go b/cmd/lambda-extension/main.go
new file mode 100644
index 0000000..056419f
--- /dev/null
+++ b/cmd/lambda-extension/main.go
@@ -0,0 +1,141 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
+	"golang.org/x/time/rate"
+
+	"github.com/atlassian/gostatsd"
+	"github.com/atlassian/gostatsd/internal/awslambda/extension"
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+	"github.com/atlassian/gostatsd/internal/util"
+	"github.com/atlassian/gostatsd/pkg/statsd"
+	"github.com/atlassian/gostatsd/pkg/transport"
+)
+
+var (
+	Version   string
+	BuildDate string
+	GitCommit string
+)
+
+const (
+	ParamConfigPath     = "config-path"
+	ParamVerbose        = "verbose"
+	ParamLambdaFileName = "lambda-entrypoint-name"
+
+	GostatsdForwarderMode = "forwarder"
+)
+
+func GetConfiguration() (*viper.Viper, error) {
+	v := viper.New()
+	util.InitViper(v, "")
+
+	cmd := pflag.NewFlagSet(os.Args[0], pflag.ContinueOnError)
+	cmd.Bool(ParamVerbose, false, "Enables debug level logs within the extension")
+	cmd.String(ParamConfigPath, "", "Path to a configuration file")
+	cmd.String(ParamLambdaFileName, "", "Name of executable that boostraps lambda. Used as Lambda-Extension-Name header")
+
+	gostatsd.AddFlags(cmd)
+
+	cmd.VisitAll(func(f *pflag.Flag) {
+		if err := v.BindPFlag(f.Name, f); err != nil {
+			panic(err)
+		}
+	})
+
+	if err := cmd.Parse(os.Args[1:]); err != nil {
+		return nil, err
+	}
+
+	if confPath := v.GetString(ParamConfigPath); confPath != "" {
+		v.SetConfigFile(confPath)
+		if err := v.ReadInConfig(); err != nil {
+			return nil, err
+		}
+	}
+
+	return v, nil
+}
+
+func CreateServer(v *viper.Viper, logger logrus.FieldLogger) (*statsd.Server, error) {
+	// HTTP client pool
+	pool := transport.NewTransportPool(logger, v)
+
+	// Create server in forwarder mode
+	return &statsd.Server{
+		InternalTags:      v.GetStringSlice(gostatsd.ParamInternalTags),
+		InternalNamespace: v.GetString(gostatsd.ParamInternalNamespace),
+		DefaultTags:       v.GetStringSlice(gostatsd.ParamDefaultTags),
+		Hostname:          gostatsd.Source(v.GetString(gostatsd.ParamHostname)),
+		FlushInterval:     v.GetDuration(gostatsd.ParamFlushInterval),
+		IgnoreHost:        v.GetBool(gostatsd.ParamIgnoreHost),
+		MaxReaders:        v.GetInt(gostatsd.ParamMaxReaders),
+		MaxParsers:        v.GetInt(gostatsd.ParamMaxParsers),
+		EstimatedTags:     v.GetInt(gostatsd.ParamEstimatedTags),
+		MetricsAddr:       v.GetString(gostatsd.ParamMetricsAddr),
+		Namespace:         v.GetString(gostatsd.ParamNamespace),
+		StatserType:       v.GetString(gostatsd.ParamStatserType),
+		HeartbeatEnabled:  v.GetBool(gostatsd.ParamHeartbeatEnabled),
+		ReceiveBatchSize:  v.GetInt(gostatsd.ParamReceiveBatchSize),
+		ConnPerReader:     v.GetBool(gostatsd.ParamConnPerReader),
+		ServerMode:        GostatsdForwarderMode,
+		LogRawMetric:      v.GetBool(gostatsd.ParamLogRawMetric),
+		HeartbeatTags: gostatsd.Tags{
+			fmt.Sprintf("version:%s", Version),
+			fmt.Sprintf("commit:%s", GitCommit),
+		},
+		BadLineRateLimitPerSecond: rate.Limit(v.GetFloat64(gostatsd.ParamBadLinesPerMinute) / 60.0),
+		Viper:                     v,
+		TransportPool:             pool,
+	}, nil
+}
+
+func main() {
+	conf, err := GetConfiguration()
+	if err != nil {
+		panic(err)
+	}
+
+	if Version == "" {
+		Version = GetVersion()
+	}
+
+	log := logrus.New().WithFields(map[string]interface{}{
+		"version":   Version,
+		"buildDate": BuildDate,
+		"gitCommit": GitCommit,
+	})
+
+	log.Logger.SetFormatter(&logrus.JSONFormatter{})
+
+	if conf.GetBool(ParamVerbose) {
+		log.Logger.SetLevel(logrus.DebugLevel)
+	}
+
+	// TODO: Configure log group
+
+	log.Info("Starting extension runtime")
+
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+	defer cancel()
+
+	manager := extension.NewManager(os.Getenv(api.EnvLambdaAPIHostname), conf.GetString(ParamLambdaFileName), log)
+
+	server, err := CreateServer(conf, log)
+	if err != nil {
+		log.WithError(err).Panic("Unable to create gostatsd server")
+	}
+
+	if err := manager.Run(ctx, server); err != nil {
+		log.WithError(err).Error("Failed trying to run lambda extension")
+	}
+
+	log.Info("Shutting down")
+}
diff --git a/cmd/lambda-extension/version.go b/cmd/lambda-extension/version.go
new file mode 100644
index 0000000..b255124
--- /dev/null
+++ b/cmd/lambda-extension/version.go
@@ -0,0 +1,10 @@
+package main
+
+import "runtime/debug"
+
+func GetVersion() string {
+	if build, ok := debug.ReadBuildInfo(); ok {
+		return build.Main.Version
+	}
+	return "unknown"
+}
diff --git a/internal/awslambda/extension/api/api.go b/internal/awslambda/extension/api/api.go
new file mode 100644
index 0000000..7e74ec5
--- /dev/null
+++ b/internal/awslambda/extension/api/api.go
@@ -0,0 +1,72 @@
+package api
+
+import (
+	"path"
+)
+
+type LambdaPath string
+
+const (
+	EnvLambdaAPIHostname = "AWS_LAMBDA_RUNTIME_API"
+	// Version defines the used OpenAPI spec being adhere to
+	// by the connecting client. It is used as part of the
+	// generated URI to be used to issue connections.
+	Version = "2020-01-01"
+
+	RegisterEndpoint  LambdaPath = "/" + Version + "/extension/register"
+	EventEndpoint     LambdaPath = "/" + Version + "/extension/event/next"
+	InitErrorEndpoint LambdaPath = "/" + Version + "/extension/init/error"
+	ExitErrorEndpoint LambdaPath = "/" + Version + "/extension/exit/error"
+
+	Invoke   Event = "INVOKE"
+	Shutdown Event = "SHUTDOWN"
+
+	LambdaExtensionNameHeaderKey       = "Lambda-Extension-Name"
+	LambdaExtensionIdentifierHeaderKey = "Lambda-Extension-Identifier"
+	LambdaErrorHeaderKey               = "Lambda-Extension-Function-Error-Type"
+	LambdaExtensionEventIdentifer      = "Lambda-Extension-Event-Identifier"
+)
+
+func (lp LambdaPath) GetUrl(domain string) string {
+	return "http://" + path.Join(domain, string(lp))
+}
+
+func (lp LambdaPath) String() string { return string(lp) }
+
+// Defining types used within the initialise phase of the extension
+type Event string
+
+type RegisterRequestPayload struct {
+	Events []Event `json:"events"`
+}
+
+type RegisterResponsePayload struct {
+	FunctionName    string `json:"functionName"`
+	FunctionVersion string `json:"functionVersion"`
+	Handler         string `json:"handler"`
+}
+
+// Defining types used within the next phase of the extension
+type EventNextPayload struct {
+	EventType          Event             `json:"eventType"`
+	Deadline           int64             `json:"deadlineMs"`
+	RequestID          string            `json:"requestId"`
+	InvokedFunctionARN string            `json:"invokedFunctionArn"`
+	Tracing            map[string]string `json:"tracing"`
+	ShutdownReason     string            `json:"shutdownReason"`
+}
+
+// Defining the types used when processing errors
+type ErrorRequest struct {
+	Message    string   `json:"errorMessage"`
+	Type       string   `json:"errorType"`
+	StackTrace []string `json:"stackTrace"`
+}
+type ErrorResponse struct {
+	Message string `json:"errorMessage"`
+	Type    string `json:"errorType"`
+}
+
+type StatusResponse struct {
+	Status string `json:"status"`
+}
diff --git a/internal/awslambda/extension/api/doc.go b/internal/awslambda/extension/api/doc.go
new file mode 100644
index 0000000..99537a8
--- /dev/null
+++ b/internal/awslambda/extension/api/doc.go
@@ -0,0 +1,8 @@
+// Package api defines the interaction that is done between
+// the aditional runtime (the current running extension)
+// and the main business logic function running along aside it.
+//
+// For further infomation regarding lambda extensions:
+//
+//	https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html#runtimes-extensions-registration-api
+package api
diff --git a/internal/awslambda/extension/manager.go b/internal/awslambda/extension/manager.go
new file mode 100644
index 0000000..60b9550
--- /dev/null
+++ b/internal/awslambda/extension/manager.go
@@ -0,0 +1,325 @@
+package extension
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"runtime/debug"
+	"strings"
+	"time"
+
+	"github.com/ash2k/stager/wait"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/sirupsen/logrus"
+	"github.com/tilinna/clock"
+	"go.uber.org/multierr"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+)
+
+const (
+	lambdaRuntimeErrorHeader = "extension.runtime.fault"
+	lambdaFailedInitHeader   = "extension.initialization.failed"
+)
+
+var (
+	ErrFailedRegistration = errors.New("extension failed to register with lambda")
+	ErrIssueProgress      = errors.New("unable to continue execution")
+	ErrServerEarlyExit    = errors.New("server has shutdown early without cause")
+)
+
+// Manager ensures that the lifetime management of the lambda
+type Manager interface {
+	Run(ctx context.Context, server Server) error
+}
+
+// Server is an interface to the gostatsd.Server type,
+// the main purpose here is to allow for mocks to be used throughout testing
+// and reduce the amount of set up code to test
+type Server interface {
+	Run(ctx context.Context) error
+}
+
+type manager struct {
+	log    logrus.FieldLogger
+	client *http.Client
+
+	domain       string
+	name         string
+	registeredID string
+}
+
+var _ Manager = (*manager)(nil)
+
+func NewManager(lambdaDomain string, lambdaFileName string, log logrus.FieldLogger) Manager {
+	m := &manager{
+		log:    log,
+		client: &http.Client{},
+		domain: lambdaDomain,
+		name:   lambdaFileName,
+	}
+
+	return m
+}
+
+func (m *manager) Run(parent context.Context, server Server) error {
+	// Init Phase of the lambda
+	if err := m.register(parent); err != nil {
+		return err
+	}
+
+	ctx, cancel := context.WithCancel(parent)
+	// Start gostatsd server
+	var wg wait.Group
+	var chErrs = make(chan error, 2)
+	defer cancel()
+
+	wg.StartWithContext(ctx, func(c context.Context) {
+		err := server.Run(c)
+		// Shutting down due to context is an acceptable
+		// reason to shutdown and is not considered an error
+		// we write all other errors to the channel
+		if !errors.Is(err, ctx.Err()) || err == nil {
+			chErrs <- err
+		}
+	})
+
+	select {
+	case err := <-chErrs:
+		// In the event that the lambda finished early before
+		// it had started accepting events without error,
+		// It is still considered an error since
+		// it was not correctly shutdown
+		if err == nil {
+			err = ErrServerEarlyExit
+		}
+		return multierr.Combine(err, m.initError(parent, err))
+	case <-clock.After(ctx, 100*time.Millisecond):
+		// In the event that statsd server returns an
+		// error during the allowed start up time
+		// the we can fail the registration
+	}
+
+	wg.StartWithContext(ctx, func(c context.Context) {
+		chErrs <- m.heartbeat(c)
+	})
+
+	wg.Wait()
+	close(chErrs)
+
+	var err error
+	for er := range chErrs {
+		err = multierr.Append(err, er)
+	}
+
+	if err != nil {
+		if errors.Is(err, ErrIssueProgress) {
+			return err
+		}
+		// Since context can be closed here, a seperate context is used
+		// to control sending exit data to the lambda
+		ctx, cancel := clock.TimeoutContext(context.Background(), time.Second)
+		defer cancel()
+
+		m.reportExitError(ctx, err)
+		return err
+	}
+
+	return nil
+}
+
+func (m *manager) heartbeat(ctx context.Context) error {
+	for ctx.Err() == nil {
+		resp, err := m.nextEvent(ctx)
+		if err != nil {
+			return err
+		}
+
+		if resp.EventType == api.Shutdown {
+			m.log.WithField("reason", resp.ShutdownReason).Info("Shutting down extension handler")
+			break
+		}
+		m.log.WithFields(map[string]interface{}{
+			"requestId":     resp.RequestID,
+			"invocationArn": resp.InvokedFunctionARN,
+			"deadline":      resp.Deadline,
+		}).Debug("Progressing further with the invocation")
+	}
+
+	return nil
+}
+
+func (m *manager) register(ctx context.Context) error {
+	var buf bytes.Buffer
+
+	err := jsoniter.NewEncoder(&buf).Encode(&api.RegisterRequestPayload{
+		Events: []api.Event{api.Invoke, api.Shutdown},
+	})
+
+	if err != nil {
+		return err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, api.RegisterEndpoint.GetUrl(m.domain), &buf)
+	if err != nil {
+		return err
+	}
+	req.Header.Set(api.LambdaExtensionNameHeaderKey, m.name)
+
+	resp, err := m.client.Do(req)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		_ = resp.Body.Close()
+	}()
+
+	switch resp.StatusCode {
+	case http.StatusOK:
+		// All things are going well so far
+	default:
+		return multierr.Combine(
+			fmt.Errorf("issue trying connect to extension with status code %d", resp.StatusCode),
+			ErrFailedRegistration,
+		)
+	}
+
+	// Once we have successfully registered to the lambda,
+	// the id assigned to the process needs to be preserved and sent
+	// with future requests
+	m.registeredID = resp.Header.Get(api.LambdaExtensionIdentifierHeaderKey)
+	if m.registeredID == "" {
+		return multierr.Combine(
+			errors.New("missing required identifier header in response"),
+			ErrFailedRegistration,
+		)
+	}
+
+	var info api.RegisterResponsePayload
+	if err := jsoniter.NewDecoder(resp.Body).Decode(&info); err != nil {
+		return err
+	}
+
+	// Log the registered payload here as an informative means of
+	// debugging connecitivity issues in future.
+	m.log.WithFields(map[string]interface{}{
+		"functionName":    info.FunctionName,
+		"functionVersion": info.FunctionVersion,
+		"functionHandler": info.Handler,
+	}).Info("Successfully registered with Lambda")
+	return nil
+}
+
+func (m *manager) nextEvent(ctx context.Context) (*api.EventNextPayload, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, api.EventEndpoint.GetUrl(m.domain), http.NoBody)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set(api.LambdaExtensionIdentifierHeaderKey, m.registeredID)
+
+	resp, err := m.client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		_ = resp.Body.Close()
+	}()
+
+	switch resp.StatusCode {
+	case http.StatusOK:
+		// All things are good in the world
+	default:
+		return nil, multierr.Combine(
+			fmt.Errorf("issue handling request with status code %d", resp.StatusCode),
+			ErrIssueProgress,
+		)
+	}
+
+	var payload api.EventNextPayload
+	if err := jsoniter.NewDecoder(resp.Body).Decode(&payload); err != nil {
+		return nil, err
+	}
+
+	return &payload, nil
+}
+
+func (m *manager) initError(ctx context.Context, problem error) error {
+	var buf bytes.Buffer
+
+	err := jsoniter.NewEncoder(&buf).Encode(api.ErrorRequest{
+		Message:    problem.Error(),
+		Type:       "InitializationError",
+		StackTrace: strings.Fields(string(debug.Stack())),
+	})
+	if err != nil {
+		return err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, api.InitErrorEndpoint.GetUrl(m.domain), &buf)
+	if err != nil {
+		return err
+	}
+
+	req.Header.Set(api.LambdaExtensionIdentifierHeaderKey, m.registeredID)
+	req.Header.Set(api.LambdaErrorHeaderKey, lambdaFailedInitHeader)
+
+	resp, err := m.client.Do(req)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		_ = resp.Body.Close()
+	}()
+
+	if resp.StatusCode == http.StatusAccepted {
+		return nil
+	}
+
+	return fmt.Errorf("issue with sending init error to lambda, status code: %d", resp.StatusCode)
+}
+
+func (m *manager) reportExitError(ctx context.Context, problem error) {
+	var buf bytes.Buffer
+
+	err := jsoniter.NewEncoder(&buf).Encode(api.ErrorRequest{
+		Message:    problem.Error(),
+		Type:       "ShutdownError",
+		StackTrace: strings.Fields(string(debug.Stack())),
+	})
+
+	if err != nil {
+		m.log.WithError(err).Error("error encoding error request to lambda runtime")
+		return
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, api.ExitErrorEndpoint.GetUrl(m.domain), &buf)
+	if err != nil {
+		m.log.WithError(err).Error("error forming http exit error request to lambda runtime")
+		return
+	}
+
+	req.Header.Set(api.LambdaExtensionIdentifierHeaderKey, m.registeredID)
+	req.Header.Set(api.LambdaErrorHeaderKey, lambdaRuntimeErrorHeader)
+
+	resp, err := m.client.Do(req)
+	if err != nil {
+		m.log.WithError(err).Error("error submitting exit error to lambda runtime")
+		return
+	}
+
+	defer func() {
+		_, _ = io.Copy(io.Discard, resp.Body)
+		_ = resp.Body.Close()
+	}()
+
+	if resp.StatusCode != http.StatusAccepted {
+		m.log.WithField("status code", resp.StatusCode).Error("issue with sending exit error to lambda")
+	}
+}
diff --git a/internal/awslambda/extension/spec/extensions-api.yaml b/internal/awslambda/extension/spec/extensions-api.yaml
new file mode 100644
index 0000000..fec80f9
--- /dev/null
+++ b/internal/awslambda/extension/spec/extensions-api.yaml
@@ -0,0 +1,320 @@
+# This document describes the AWS Lambda Runtime Extensions API using the OpenAPI 3.0 specification.
+
+openapi: 3.0.0
+info:
+  title: AWS Lambda Runtime Extensions API
+  description: AWS Lambda Runtime Extensions API is an HTTP API for implementing runtime extensions
+  version: 1.0.1
+
+servers:
+  - url: https://example.com/2020-01-01
+
+paths:
+
+  /extension/register:
+    post:
+      summary: >
+        Register an extension with the given name.
+      parameters:
+        - in: header
+          name: Lambda-Extension-Name
+          required: true
+          description: Public extension name.
+          schema:
+            type: string
+      requestBody:
+        content:
+          application/vnd.aws.lambda.extension+json:
+            schema:
+              $ref: '#/components/schemas/RegisterRequest'
+      responses:
+        '200':
+          description: >
+            Response contains generated unique extension identifier that is required for all subsequent requests.
+          headers:
+            Lambda-Extension-Identifier:
+              description: Generated unique identifier for public extension name.
+              schema:
+                type: string
+          content:
+            application/vnd.aws.lambda.extension+json:
+              schema:
+                $ref: '#/components/schemas/RegisterResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Extension should exit promptly.
+
+  /extension/event/next:
+    get:
+      summary: >
+        Extension makes this HTTP request when it is ready to receive and process a new event.
+      parameters:
+        - in: header
+          name: Lambda-Extension-Identifier
+          required: true
+          description: Unique identifier for extension.
+          schema:
+            type: string
+      responses:
+        '200':
+          description: >
+            This is an iterator-style blocking API call. Response contains
+            event JSON document.
+          headers:
+            Lambda-Extension-Event-Identifier:
+              description: Event UUID.
+              schema:
+                type: string
+          content:
+            application/vnd.aws.lambda.extensionevent+json:
+              schema:
+                oneOf:
+                - $ref: '#/components/schemas/EventInvoke'
+                - $ref: '#/components/schemas/EventShutdown'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Extension should exit promptly.
+
+  /extension/init/error:
+    post:
+      summary: >
+        This allows the extension to report an initialization error to the platform.
+        Call it when you registered but failed to initialize. Once error has been
+        accepted by the platform, no further API calls will succeed, extension should
+        exit.
+      parameters:
+        - in: header
+          name: Lambda-Extension-Identifier
+          required: true
+          description: Unique identifier for extension.
+          schema:
+            type: string
+        - in: header
+          name: Lambda-Extension-Function-Error-Type
+          required: true
+          description: Extension.ErrorEnum
+          schema:
+            type: string
+      requestBody:
+        content:
+          application/vnd.aws.lambda.error+json:
+            schema:
+              $ref: '#/components/schemas/ErrorRequest'
+      responses:
+        '202':
+          description: Accepted
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/StatusResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Extension should exit promptly.
+
+  /extension/exit/error:
+    post:
+      summary: >
+        This allows the extension to report an error to the platform before exiting.
+        Call it when you encounter an unexpected failure. Once error has been
+        accepted by the platform, no further API calls will succeed, extension should
+        exit.
+      parameters:
+        - in: header
+          name: Lambda-Extension-Identifier
+          required: true
+          description: Unique identifier for extension.
+          schema:
+            type: string
+        - in: header
+          name: Lambda-Extension-Function-Error-Type
+          required: true
+          description: Extension.ErrorEnum
+          schema:
+            type: string
+      requestBody:
+        content:
+          application/vnd.aws.lambda.error+json:
+            schema:
+              $ref: '#/components/schemas/ErrorRequest'
+      responses:
+        '202':
+          description: Accepted
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/StatusResponse'
+        '400':
+          description: Bad Request
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '403':
+          description: Forbidden
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/ErrorResponse'
+        '500':
+          description: >
+            Container error. Non-recoverable state. Extension should exit promptly.
+
+components:
+  schemas:
+    RegisterRequest:
+      type: object
+      properties:
+        events:
+          type: array
+          items:
+            type: string
+      example:
+        {
+          'events': [ 'INVOKE', 'SHUTDOWN']
+        }
+
+    RegisterResponse:
+      type: object
+      properties:
+        functionName:
+          type: string
+        functionVersion:
+          type: string
+        handler:
+          type: string
+      example:
+        {
+          functionName: 'my-function',
+          functionVersion: '$LATEST',
+          handler: 'lambda_handler'
+        }
+
+    ErrorResponse:
+      type: object
+      properties:
+        errorMessage:
+          type: string
+        errorType:
+          type: string
+      example:
+        {
+          errorMessage: 'human readable error message',
+          errorType: 'Platform.ErrorEnum',
+        }
+
+    ErrorRequest:
+      type: object
+      properties:
+        errorMessage:
+          type: string
+        errorType:
+          type: string
+        stackTrace:
+          type: array
+          items:
+            type: string
+      example:
+        {
+          errorType: 'Extension.ErrorEnum',
+          errorMessage: 'human readable error message',
+          stackTrace: []
+        }
+
+    EventInvoke:
+      type: object
+      properties:
+        eventType:
+          type: string
+        deadlineMs:
+          type: integer
+          description: The date that the function times out in Unix time milliseconds.
+        requestId:
+          type: string
+        invokedFunctionArn:
+          type: string
+        tracing:
+          oneOf:
+            - $ref: '#/components/schemas/XRayTracingInfo'
+      example:
+        {
+          eventType: 'INVOKE',
+          deadlineMs: 1581512138111,
+          requestId: 'aws-request-ID',
+          invokedFunctionArn: 'invoked-function-arn',
+          tracing: {
+            type: 'X-Amzn-Trace-Id',
+            value: 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1'
+          }
+        }
+
+    EventShutdown:
+      type: object
+      properties:
+        eventType:
+          type: string
+        shutdownReason:
+          type: string
+          enum: [spindown, timeout, failure]
+        deadlineMs:
+          type: integer
+          description: The date that the function times out in Unix time milliseconds.
+      example:
+        {
+          eventType: 'SHUTDOWN',
+          shutdownReason: 'spindown',
+          deadlineMs: 1581512138111
+        }
+
+    StatusResponse:
+      type: object
+      properties:
+        status:
+          type: string
+      example:
+        {
+          status: "Accepted"
+        }
+
+    XRayTracingInfo:
+      type: object
+      properties:
+        type:
+          type: string
+          enum: [X-Amzn-Trace-Id]
+        value:
+          type: string
+      example:
+        {
+          type: 'X-Amzn-Trace-Id',
+          value: 'Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1'
+        }


Test patch:
diff --git a/internal/awslambda/extension/api/api_test.go b/internal/awslambda/extension/api/api_test.go
new file mode 100644
index 0000000..7b55550
--- /dev/null
+++ b/internal/awslambda/extension/api/api_test.go
@@ -0,0 +1,26 @@
+package api_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+)
+
+func TestExpectedURL(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		domain string
+		path   api.LambdaPath
+		expect string
+	}{
+		{domain: "localhost", path: api.EventEndpoint, expect: "http://localhost/2020-01-01/extension/event/next"},
+		{domain: "localhost/", path: api.ExitErrorEndpoint, expect: "http://localhost/2020-01-01/extension/exit/error"},
+	}
+
+	for _, tc := range testCases {
+		assert.Equal(t, tc.expect, tc.path.GetUrl(tc.domain))
+	}
+}
diff --git a/internal/awslambda/extension/manager_test.go b/internal/awslambda/extension/manager_test.go
new file mode 100644
index 0000000..da7869b
--- /dev/null
+++ b/internal/awslambda/extension/manager_test.go
@@ -0,0 +1,337 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/gorilla/mux"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tilinna/clock"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+	"github.com/atlassian/gostatsd/internal/fixtures"
+	"github.com/atlassian/gostatsd/pkg/fakesocket"
+)
+
+type mocked struct {
+	delay time.Duration
+	erred error
+
+	start chan<- struct{}
+	done  chan<- struct{}
+}
+
+func (m *mocked) Run(ctx context.Context) error {
+	m.start <- struct{}{}
+	if m.delay > 0 {
+		<-clock.After(ctx, m.delay)
+	}
+	m.done <- struct{}{}
+	return m.erred
+}
+
+func InitHandler(tb testing.TB, statusCode int) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		// Enforcing correct http method is used
+		require.Equal(tb, http.MethodPost, r.Method)
+		// Enforcing Incoming headers as per the spec
+		require.Contains(tb, r.Header, api.LambdaExtensionNameHeaderKey)
+		// Encorcing that payload is correct
+		var payload api.RegisterRequestPayload
+		assert.NoError(tb, json.NewDecoder(r.Body).Decode(&payload), "Must send a valid payload request")
+
+		switch statusCode {
+		case http.StatusOK:
+			rw.Header().Set(api.LambdaExtensionIdentifierHeaderKey, tb.Name())
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.RegisterResponsePayload{
+				FunctionName:    "mock-init-handler",
+				FunctionVersion: "1.0.0-mock",
+				Handler:         tb.Name(),
+			}))
+		case http.StatusBadRequest, http.StatusForbidden:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.ErrorResponse{
+				Message: "Error failed successfully",
+				Type:    "mock-test.error",
+			}))
+		default:
+			rw.WriteHeader(statusCode)
+		}
+	})
+}
+
+func InitErrorHandler(tb testing.TB, statusCode int) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		// Enforcing that the correct method is used
+		require.Equal(tb, http.MethodPost, r.Method)
+		// Enforcing that the headers exist
+		require.Contains(tb, r.Header, api.LambdaExtensionIdentifierHeaderKey)
+		require.Contains(tb, r.Header, api.LambdaErrorHeaderKey)
+
+		var payload api.ErrorRequest
+		assert.NoError(tb, json.NewDecoder(r.Body).Decode(&payload))
+
+		switch statusCode {
+		case http.StatusAccepted:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.StatusResponse{
+				Status: "accepted",
+			}))
+		case http.StatusBadRequest, http.StatusForbidden:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.ErrorResponse{
+				Message: "Unable to process init errored shutdown",
+				Type:    "mock-test.error",
+			}))
+		default:
+			rw.WriteHeader(statusCode)
+		}
+	})
+}
+
+func ExitErrorHandler(tb testing.TB, statusCode int) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		// Enforcing that the correct method is used
+		require.Equal(tb, http.MethodPost, r.Method)
+		// Enforcing that the headers exist
+		require.Contains(tb, r.Header, api.LambdaExtensionIdentifierHeaderKey)
+		require.Contains(tb, r.Header, api.LambdaErrorHeaderKey)
+
+		var payload api.ErrorRequest
+		assert.NoError(tb, json.NewDecoder(r.Body).Decode(&payload))
+
+		switch statusCode {
+		case http.StatusAccepted:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.StatusResponse{
+				Status: "accepted",
+			}))
+		case http.StatusBadRequest, http.StatusForbidden:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.ErrorResponse{
+				Message: "Unable to process exit errored shutdown",
+				Type:    "mock-test.error",
+			}))
+		default:
+			rw.WriteHeader(statusCode)
+		}
+	})
+}
+
+func EventNextHandler(tb testing.TB, ctx context.Context, statusCode, shutdownAfter int) http.Handler {
+	count := 0
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		// Enforcing the correct method is used
+		require.Equal(tb, http.MethodGet, r.Method)
+		// Enforcing the header exist
+		require.Contains(tb, r.Header, api.LambdaExtensionIdentifierHeaderKey)
+
+		switch statusCode {
+		case http.StatusOK:
+			payload := &api.EventNextPayload{
+				EventType:          api.Invoke,
+				Deadline:           clock.FromContext(ctx).Now().UTC().Add(100 * time.Millisecond).UnixMilli(),
+				RequestID:          fmt.Sprint(count),
+				InvokedFunctionARN: "local:dev:gostatsd-extension-manager",
+			}
+			if count > shutdownAfter {
+				payload.EventType = api.Shutdown
+				payload.ShutdownReason = "Finished running test cases"
+			}
+			assert.NoError(tb, json.NewEncoder(rw).Encode(payload))
+			rw.Header().Set(api.LambdaExtensionEventIdentifer, fmt.Sprint(count))
+		case http.StatusForbidden:
+			rw.WriteHeader(statusCode)
+			assert.NoError(tb, json.NewEncoder(rw).Encode(&api.ErrorResponse{
+				Message: "Failed to process next handler",
+				Type:    "mock-test.error",
+			}))
+		default:
+			rw.WriteHeader(statusCode)
+		}
+
+		count++
+	})
+}
+
+func TestManagerRegister(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Scenario   string
+		StatusCode int
+		Expect     error
+	}{
+		{
+			Scenario:   "Registered to Lambda",
+			StatusCode: http.StatusOK,
+			Expect:     nil,
+		},
+		{
+			Scenario:   "Invalid register payload sent",
+			StatusCode: http.StatusBadRequest,
+			Expect:     ErrFailedRegistration,
+		},
+		{
+			Scenario:   "Forbidden request sent to Lambda",
+			StatusCode: http.StatusForbidden,
+			Expect:     ErrFailedRegistration,
+		},
+		{
+			Scenario:   "Unknown error within lambda",
+			StatusCode: http.StatusInternalServerError,
+			Expect:     ErrFailedRegistration,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.Scenario, func(t *testing.T) {
+			t.Parallel()
+			r := mux.NewRouter()
+			r.Handle(api.RegisterEndpoint.String(), InitHandler(t, tc.StatusCode))
+			r.PathPrefix("/").HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+				assert.Fail(t, "Trying to access wrong path", r.RequestURI)
+			})
+
+			s := httptest.NewServer(r)
+			t.Cleanup(s.Close)
+
+			u, err := url.Parse(s.URL)
+			require.NoError(t, err, "Must be able to parse URL from httptest server")
+
+			m := &manager{
+				log:    logrus.New().WithField("test-case", t.Name()),
+				client: s.Client(),
+				domain: u.Hostname() + ":" + u.Port(),
+				name:   t.Name(),
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			t.Cleanup(cancel)
+
+			assert.ErrorIs(t, m.register(ctx), tc.Expect)
+		})
+	}
+}
+
+func TestManagerDo(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Scenario      string
+		EventStatus   int
+		ShutdownAfter int
+		MockDelay     time.Duration
+		MockError     error
+
+		ExpectError error
+	}{
+		{
+			Scenario:      "Normal operation",
+			EventStatus:   http.StatusOK,
+			ShutdownAfter: 3,
+			MockDelay:     time.Second,
+			MockError:     nil,
+			ExpectError:   nil,
+		},
+		{
+			Scenario:      "Operational failure sending heartbeart",
+			EventStatus:   http.StatusInternalServerError,
+			ShutdownAfter: 2,
+			MockDelay:     time.Second,
+			MockError:     nil,
+			ExpectError:   ErrIssueProgress,
+		},
+		{
+			Scenario:      "Server has shutdown early without cause",
+			EventStatus:   http.StatusOK,
+			ShutdownAfter: 3,
+			MockDelay:     0,
+			MockError:     nil,
+			ExpectError:   ErrServerEarlyExit,
+		},
+		{
+			Scenario:      "Server configuration was wrong resulting in during init",
+			EventStatus:   http.StatusOK,
+			ShutdownAfter: 0,
+			MockDelay:     0,
+			MockError:     fakesocket.ErrClosedConnection,
+			ExpectError:   fakesocket.ErrClosedConnection,
+		},
+		{
+			Scenario:      "Server failed throughout runtime and successfully committed to lambda",
+			EventStatus:   http.StatusOK,
+			ShutdownAfter: 0,
+			MockDelay:     300 * time.Millisecond,
+			MockError:     fakesocket.ErrClosedConnection,
+			ExpectError:   fakesocket.ErrClosedConnection,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.Scenario, func(t *testing.T) {
+			t.Parallel()
+			clck := clock.NewMock(time.Unix(1, 0))
+
+			ctx, cancel := context.WithCancel(
+				clock.Context(context.Background(), clck),
+			)
+			t.Cleanup(cancel)
+
+			r := mux.NewRouter()
+			r.Handle(api.RegisterEndpoint.String(), InitHandler(t, http.StatusOK))
+			r.Handle(api.InitErrorEndpoint.String(), InitErrorHandler(t, http.StatusAccepted))
+			r.Handle(api.ExitErrorEndpoint.String(), ExitErrorHandler(t, http.StatusAccepted))
+			r.Handle(api.EventEndpoint.String(), EventNextHandler(t, ctx, tc.EventStatus, tc.ShutdownAfter))
+			r.PathPrefix("/").HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+				assert.Fail(t, "Trying to access wrong path", r.RequestURI)
+			})
+
+			s := httptest.NewServer(r)
+			t.Cleanup(s.Close)
+
+			u, err := url.Parse(s.URL)
+			require.NoError(t, err, "Must be able to parse URL from httptest server")
+
+			m := &manager{
+				log:    logrus.New(),
+				client: s.Client(),
+				domain: u.Hostname() + ":" + u.Port(),
+				name:   t.Name(),
+			}
+
+			start := make(chan struct{}, 1)
+			done := make(chan struct{}, 1)
+
+			server := &mocked{delay: tc.MockDelay, erred: tc.MockError, start: start, done: done}
+
+			go func() {
+				<-start
+				close(start)
+				for {
+					select {
+					case <-done:
+						close(done)
+						fixtures.NextStep(ctx, clck)
+						cancel()
+						return
+					default:
+						fixtures.NextStep(ctx, clck)
+					}
+				}
+			}()
+			assert.ErrorIs(t, m.Run(ctx, server), tc.ExpectError)
+			assert.Zero(t, clck.Len(), "Must have stopped all timers")
+		})
+	}
+}