Add capability to define optional handler behavior

lambda/entry.go

@@ -37,12 +37,14 @@ import (
 // Where "TIn" and "TOut" are types compatible with the "encoding/json" standard library.
 // See https://golang.org/pkg/encoding/json/#Unmarshal for how deserialization behaves
 func Start(handler interface{}) {
-	StartWithContext(context.Background(), handler)
+	StartWithOptions(handler)
 }
 
 // StartWithContext is the same as Start except sets the base context for the function.
+//
+// Deprecated: use lambda.StartWithOptions(handler, lambda.WithContext(ctx)) instead
 func StartWithContext(ctx context.Context, handler interface{}) {
-	StartHandlerWithContext(ctx, NewHandler(handler))
+	StartWithOptions(handler, WithContext(ctx))
 }
 
 // StartHandler takes in a Handler wrapper interface which can be implemented either by a
@@ -51,13 +53,20 @@ func StartWithContext(ctx context.Context, handler interface{}) {
 // Handler implementation requires a single "Invoke()" function:
 //
 //  func Invoke(context.Context, []byte) ([]byte, error)
+//
+// Deprecated: use lambda.Start(handler) instead
 func StartHandler(handler Handler) {
-	StartHandlerWithContext(context.Background(), handler)
+	StartWithOptions(handler)
+}
+
+// StartWithOptions is the same as Start after the application of any handler options specified
+func StartWithOptions(handler interface{}, options ...Option) {
+	start(newHandler(handler, options...))
 }
 
 type startFunction struct {
 	env string
-	f   func(ctx context.Context, envValue string, handler Handler) error
+	f   func(envValue string, handler Handler) error
 }
 
 var (
@@ -66,7 +75,7 @@ var (
 	// To drop the rpc dependencies, compile with `-tags lambda.norpc`
 	rpcStartFunction = &startFunction{
 		env: "_LAMBDA_SERVER_PORT",
-		f: func(c context.Context, p string, h Handler) error {
+		f: func(_ string, _ Handler) error {
 			return errors.New("_LAMBDA_SERVER_PORT was present but the function was compiled without RPC support")
 		},
 	}
@@ -85,17 +94,24 @@ var (
 // Handler implementation requires a single "Invoke()" function:
 //
 //  func Invoke(context.Context, []byte) ([]byte, error)
+//
+// Deprecated: use lambda.StartWithOptions(handler, lambda.WithContext(ctx)) instead
 func StartHandlerWithContext(ctx context.Context, handler Handler) {
+	StartWithOptions(handler, WithContext(ctx))
+}
+
+func start(handler *handlerOptions) {
 	var keys []string
 	for _, start := range startFunctions {
 		config := os.Getenv(start.env)
 		if config != "" {
 			// in normal operation, the start function never returns
 			// if it does, exit!, this triggers a restart of the lambda function
-			err := start.f(ctx, config, handler)
+			err := start.f(config, handler)
 			logFatalf("%v", err)
 		}
 		keys = append(keys, start.env)
 	}
 	logFatalf("expected AWS Lambda environment variables %s are not defined", keys)
+
 }

lambda/function.go

@@ -13,14 +13,17 @@ import (
 )
 
 // Function struct which wrap the Handler
+//
+// Deprecated: The Function type is public for the go1.x runtime internal use of the net/rpc package
 type Function struct {
-	handler Handler
-	ctx     context.Context
+	handler *handlerOptions
 }
 
 // NewFunction which creates a Function with a given Handler
+//
+// Deprecated: The Function type is public for the go1.x runtime internal use of the net/rpc package
 func NewFunction(handler Handler) *Function {
-	return &Function{handler: handler}
+	return &Function{newHandler(handler)}
 }
 
 // Ping method which given a PingRequest and a PingResponse parses the PingResponse
@@ -38,7 +41,7 @@ func (fn *Function) Invoke(req *messages.InvokeRequest, response *messages.Invok
 	}()
 
 	deadline := time.Unix(req.Deadline.Seconds, req.Deadline.Nanos).UTC()
-	invokeContext, cancel := context.WithDeadline(fn.context(), deadline)
+	invokeContext, cancel := context.WithDeadline(fn.baseContext(), deadline)
 	defer cancel()
 
 	lc := &lambdacontext.LambdaContext{
@@ -70,26 +73,9 @@ func (fn *Function) Invoke(req *messages.InvokeRequest, response *messages.Invok
 	return nil
 }
 
-// context returns the base context used for the fn.
-func (fn *Function) context() context.Context {
-	if fn.ctx == nil {
-		return context.Background()
+func (fn *Function) baseContext() context.Context {
+	if fn.handler.baseContext != nil {
+		return fn.handler.baseContext
 	}
-
-	return fn.ctx
-}
-
-// withContext returns a shallow copy of Function with its context changed
-// to the provided ctx. If the provided ctx is non-nil a Background context is set.
-func (fn *Function) withContext(ctx context.Context) *Function {
-	if ctx == nil {
-		ctx = context.Background()
-	}
-
-	fn2 := new(Function)
-	*fn2 = *fn
-
-	fn2.ctx = ctx
-
-	return fn2
+	return context.Background()
 }

lambda/function_test.go

@@ -36,14 +36,14 @@ func (h testWrapperHandler) Invoke(ctx context.Context, payload []byte) ([]byte,
 var _ Handler = (testWrapperHandler)(nil)
 
 func TestInvoke(t *testing.T) {
-	srv := &Function{handler: testWrapperHandler(
+	srv := NewFunction(testWrapperHandler(
 		func(ctx context.Context, input []byte) (interface{}, error) {
 			if deadline, ok := ctx.Deadline(); ok {
 				return deadline.UnixNano(), nil
 			}
 			return nil, errors.New("!?!?!?!?!")
 		},
-	)}
+	))
 	deadline := time.Now()
 	var response messages.InvokeResponse
 	err := srv.Invoke(&messages.InvokeRequest{
@@ -59,15 +59,17 @@ func TestInvoke(t *testing.T) {
 
 func TestInvokeWithContext(t *testing.T) {
 	key := struct{}{}
-	srv := NewFunction(testWrapperHandler(
-		func(ctx context.Context, input []byte) (interface{}, error) {
-			assert.Equal(t, "dummy", ctx.Value(key))
-			if deadline, ok := ctx.Deadline(); ok {
-				return deadline.UnixNano(), nil
-			}
-			return nil, errors.New("!?!?!?!?!")
-		}))
-	srv = srv.withContext(context.WithValue(context.Background(), key, "dummy"))
+	srv := NewFunction(&handlerOptions{
+		Handler: testWrapperHandler(
+			func(ctx context.Context, input []byte) (interface{}, error) {
+				assert.Equal(t, "dummy", ctx.Value(key))
+				if deadline, ok := ctx.Deadline(); ok {
+					return deadline.UnixNano(), nil
+				}
+				return nil, errors.New("!?!?!?!?!")
+			}),
+		baseContext: context.WithValue(context.Background(), key, "dummy"),
+	})
 	deadline := time.Now()
 	var response messages.InvokeResponse
 	err := srv.Invoke(&messages.InvokeRequest{
@@ -86,12 +88,11 @@ type CustomError struct{}
 func (e CustomError) Error() string { return "Something bad happened!" }
 
 func TestCustomError(t *testing.T) {
-
-	srv := &Function{handler: testWrapperHandler(
+	srv := NewFunction(testWrapperHandler(
 		func(ctx context.Context, input []byte) (interface{}, error) {
 			return nil, CustomError{}
 		},
-	)}
+	))
 	var response messages.InvokeResponse
 	err := srv.Invoke(&messages.InvokeRequest{}, &response)
 	assert.NoError(t, err)
@@ -106,11 +107,11 @@ func (e *CustomError2) Error() string { return "Something bad happened!" }
 
 func TestCustomErrorRef(t *testing.T) {
 
-	srv := &Function{handler: testWrapperHandler(
+	srv := NewFunction(testWrapperHandler(
 		func(ctx context.Context, input []byte) (interface{}, error) {
 			return nil, &CustomError2{}
 		},
-	)}
+	))
 	var response messages.InvokeResponse
 	err := srv.Invoke(&messages.InvokeRequest{}, &response)
 	assert.NoError(t, err)
@@ -120,12 +121,12 @@ func TestCustomErrorRef(t *testing.T) {
 }
 
 func TestContextPlumbing(t *testing.T) {
-	srv := &Function{handler: testWrapperHandler(
+	srv := NewFunction(testWrapperHandler(
 		func(ctx context.Context, input []byte) (interface{}, error) {
 			lc, _ := lambdacontext.FromContext(ctx)
 			return lc, nil
 		},
-	)}
+	))
 	var response messages.InvokeResponse
 	err := srv.Invoke(&messages.InvokeRequest{
 		CognitoIdentityId:     "dummyident",
@@ -172,14 +173,14 @@ func TestXAmznTraceID(t *testing.T) {
 		Ctx string
 	}
 
-	srv := &Function{handler: testWrapperHandler(
+	srv := NewFunction(testWrapperHandler(
 		func(ctx context.Context, input []byte) (interface{}, error) {
 			return &XRayResponse{
 				Env: os.Getenv("_X_AMZN_TRACE_ID"),
 				Ctx: ctx.Value("x-amzn-trace-id").(string),
 			}, nil
 		},
-	)}
+	))
 
 	sequence := []struct {
 		Input    string