Merge pull request #682 from atlassian/msg/api-exposure

Feature: Expose Lambda Extension package

Author: MovieStoreGuy

pkg/lambda/extension.go

@@ -0,0 +1,40 @@
+package lambda
+
+import (
+	"github.com/sirupsen/logrus"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension"
+	"github.com/atlassian/gostatsd/internal/flush"
+	"github.com/atlassian/gostatsd/pkg/statsd"
+)
+
+// Extension exposes the internal type definition.
+//
+// Note: This is not ideal but simplifies exposing the API.
+type Extension = extension.Server
+
+// NewExtension extends a `statsd.Server` by adding a server that integrates with
+// the AWS Lambda Extension Runtime API to allow near native support for on host collection.
+// The provided server is then run within the extension in additional to the server for AWS integration.
+func NewExtension(logger logrus.FieldLogger, server *statsd.Server, opts Options) (Extension, error) {
+	if err := opts.Validate(); err != nil {
+		return nil, err
+	}
+
+	s := *server
+
+	var extOpts []extension.ManagerOpt
+	if opts.EnableManualFlush {
+		fc := flush.NewFlushCoordinator()
+		s.ForwarderFlushCoordinator = fc
+		extOpts = append(extOpts, extension.WithManualFlushEnabled(fc, opts.TelemetryAddr))
+	}
+
+	return extension.NewManager(
+		opts.RuntimeAPI,
+		opts.ExecutableName,
+		logger,
+		&s,
+		extOpts...,
+	), nil
+}

pkg/lambda/extension_test.go

@@ -0,0 +1,58 @@
+package lambda
+
+import (
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/atlassian/gostatsd/pkg/statsd"
+)
+
+func TestNewExtension(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name    string
+		options Options
+		errVal  string
+	}{
+		{
+			name:    "no options value defined",
+			options: Options{},
+			errVal:  "missing `RuntimeAPI` value; missing `ExecutableName` value",
+		},
+		{
+			name: "valid options provided",
+			options: Options{
+				RuntimeAPI:     "runtime-api",
+				ExecutableName: "bin",
+			},
+			errVal: "",
+		},
+		{
+			name: "valid options provided with enable manual flush",
+			options: Options{
+				RuntimeAPI:        "runtime-api",
+				ExecutableName:    "bin",
+				EnableManualFlush: true,
+				TelemetryAddr:     ":0",
+			},
+			errVal: "",
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			actual, err := NewExtension(logrus.New(), &statsd.Server{}, tc.options)
+			if tc.errVal != "" {
+				assert.Nil(t, actual, "Must not be a valid extension")
+				assert.EqualError(t, err, tc.errVal)
+			} else {
+				assert.NotNil(t, actual, "Must be a valid extension")
+				assert.NoError(t, err, "Must not error")
+			}
+		})
+	}
+}

pkg/lambda/options.go

@@ -0,0 +1,91 @@
+package lambda
+
+import (
+	"errors"
+	"os"
+	"path"
+
+	"github.com/spf13/pflag"
+	"go.uber.org/multierr"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+)
+
+const (
+	ParamRuntimeAPI        = `lambda-runtime-api`
+	ParamEntryPointName    = `lamda-entrypoint-name`
+	ParamEnableManualFlush = `enable-manual-flush`
+	ParamTelemetryAddr     = `telemetry-addr`
+)
+
+// Options is used to provide overrides when calling `NewExtension`
+type Options struct {
+	// RuntimeAPI is AWS service that is responsible for orchestrating
+	// the extension with the other Lambda's that are part of the deployment.
+	// This will set to `env:AWS_LAMBDA_RUNTIME_API` by default as per the docs,
+	// but can be overriden for testing purposes and validation.
+	RuntimeAPI string
+	// ExecutableName is the full file name of the lambda extension that is
+	// used to validate the bootstrap process.
+	ExecutableName string
+	// EnableManualFlush allows for Lambda's
+	// to control when the forwarder would send the accumlated metrics.
+	EnableManualFlush bool
+	// TelemetryAddr is used by the AWS runtime to publish
+	// events back to service.
+	TelemetryAddr string
+}
+
+func NewOptionsFromEnvironment() Options {
+	return Options{
+		RuntimeAPI:        os.Getenv(api.EnvLambdaAPIHostname),
+		ExecutableName:    path.Base(os.Args[0]),
+		EnableManualFlush: false,
+		TelemetryAddr:     "http://sandbox:8083",
+	}
+}
+
+// AddFlags is used to add preconfigured entries
+// into an existing `FlagSet`
+func (o *Options) AddFlags(fs *pflag.FlagSet) {
+	fs.StringVar(
+		&o.RuntimeAPI,
+		ParamRuntimeAPI,
+		o.RuntimeAPI,
+		"Sets the runtime api that is used to contact the lambda runtime service",
+	)
+	fs.StringVar(
+		&o.ExecutableName,
+		ParamEntryPointName,
+		o.ExecutableName,
+		"Sets the name of the executable name used within the bootstrap process.",
+	)
+	fs.BoolVar(
+		&o.EnableManualFlush,
+		ParamEnableManualFlush,
+		o.EnableManualFlush,
+		"When set, enables lambda(s) to force the forwarder to send data. Useful in low volume lambdas",
+	)
+	fs.StringVar(
+		&o.TelemetryAddr,
+		ParamTelemetryAddr,
+		o.TelemetryAddr,
+		"Allows for callbacks to force manual flushing of accumulated metrics",
+	)
+}
+
+// Validate ensure all the values are valid,
+// any values that not are reported as errors.
+// All invalid values are reported together.
+func (o Options) Validate() (errs error) {
+	if o.RuntimeAPI == "" {
+		errs = multierr.Append(errs, errors.New("missing `RuntimeAPI` value"))
+	}
+	if o.ExecutableName == "" {
+		errs = multierr.Append(errs, errors.New("missing `ExecutableName` value"))
+	}
+	if o.EnableManualFlush && o.TelemetryAddr == "" {
+		errs = multierr.Append(errs, errors.New("missing `TelemetryAddr` when `EnableManualFlush` is enabled"))
+	}
+	return errs
+}

pkg/lambda/options_test.go

@@ -0,0 +1,97 @@
+package lambda
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/atlassian/gostatsd/internal/awslambda/extension/api"
+)
+
+func TestDefaultOptions(t *testing.T) {
+	t.Setenv(api.EnvLambdaAPIHostname, "example")
+
+	expected := Options{
+		RuntimeAPI:        "example",
+		ExecutableName:    path.Base(os.Args[0]),
+		EnableManualFlush: false,
+		TelemetryAddr:     "http://sandbox:8083",
+	}
+
+	assert.Equal(t, expected, NewOptionsFromEnvironment(), "Must match the expected value")
+}
+
+func TestOptionsAddFlag(t *testing.T) {
+	t.Parallel()
+
+	o := &Options{
+		RuntimeAPI:        "my-awesome-api",
+		ExecutableName:    "bin",
+		EnableManualFlush: true,
+		TelemetryAddr:     ":8089",
+	}
+
+	fs := pflag.NewFlagSet(t.Name(), pflag.ContinueOnError)
+	o.AddFlags(fs)
+
+	for _, flag := range []string{
+		ParamRuntimeAPI,
+		ParamEntryPointName,
+		ParamEnableManualFlush,
+		ParamTelemetryAddr,
+	} {
+		assert.NotNil(
+			t,
+			fs.Lookup(flag),
+			"Must have a valid entry for expected flag name %q",
+			flag,
+		)
+	}
+
+}
+
+func TestOptionsValidate(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name    string
+		options Options
+		errVal  string
+	}{
+		{
+			name:    "empty options",
+			options: Options{},
+			errVal:  "missing `RuntimeAPI` value; missing `ExecutableName` value",
+		},
+		{
+			name: "enabled manual flushes",
+			options: Options{
+				EnableManualFlush: true,
+			},
+			errVal: "missing `RuntimeAPI` value; missing `ExecutableName` value; missing `TelemetryAddr` when `EnableManualFlush` is enabled",
+		},
+		{
+			name: "(simulated) Default options",
+			options: Options{
+				RuntimeAPI:     "runtime-api",
+				ExecutableName: "bin",
+			},
+			errVal: "",
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := tc.options.Validate()
+			if tc.errVal != "" {
+				assert.EqualError(t, err, tc.errVal, "Must match the expected error message")
+			} else {
+				assert.NoError(t, err, "Must not error")
+			}
+		})
+	}
+}