Support setting an optional base context for functions.

lambda/entry.go

@@ -3,6 +3,7 @@
 package lambda
 
 import (
+	"context"
 	"log"
 	"net"
 	"net/rpc"
@@ -37,8 +38,12 @@ 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{}) {
-	wrappedHandler := NewHandler(handler)
-	StartHandler(wrappedHandler)
+	StartWithContext(context.Background(), handler)
+}
+
+// StartWithContext is the same as Start except sets the base context for the function.
+func StartWithContext(ctx context.Context, handler interface{}) {
+	StartHandlerWithContext(ctx, NewHandler(handler))
 }
 
 // StartHandler takes in a Handler wrapper interface which can be implemented either by a
@@ -48,12 +53,21 @@ func Start(handler interface{}) {
 //
 //  func Invoke(context.Context, []byte) ([]byte, error)
 func StartHandler(handler Handler) {
+	StartHandlerWithContext(context.Background(), handler)
+}
+
+// StartHandlerWithContext is the same as StartHandler except sets the base context for the function.
+//
+// Handler implementation requires a single "Invoke()" function:
+//
+//  func Invoke(context.Context, []byte) ([]byte, error)
+func StartHandlerWithContext(ctx context.Context, handler Handler) {
 	port := os.Getenv("_LAMBDA_SERVER_PORT")
 	lis, err := net.Listen("tcp", "localhost:"+port)
 	if err != nil {
 		log.Fatal(err)
 	}
-	err = rpc.Register(NewFunction(handler))
+	err = rpc.Register(NewFunctionWithContext(ctx, handler))
 	if err != nil {
 		log.Fatal("failed to register handler function")
 	}

lambda/function.go

@@ -16,13 +16,31 @@ import (
 // Function struct which wrap the Handler
 type Function struct {
 	handler Handler
+	context context.Context
 }
 
 // NewFunction which creates a Function with a given Handler
 func NewFunction(handler Handler) *Function {
 	return &Function{handler: handler}
 }
 
+// NewFunctionWithContext which creates a Function with a given Handler and sets the base Context.
+func NewFunctionWithContext(ctx context.Context, handler Handler) *Function {
+	return &Function{
+		context: ctx,
+		handler: handler,
+	}
+}
+
+// Context returns the base context used for the fn.
+func (fn *Function) Context() context.Context {
+	if fn.context == nil {
+		return context.Background()
+	}
+
+	return fn.context
+}
+
 // Ping method which given a PingRequest and a PingResponse parses the PingResponse
 func (fn *Function) Ping(req *messages.PingRequest, response *messages.PingResponse) error {
 	*response = messages.PingResponse{}
@@ -44,7 +62,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(context.Background(), deadline)
+	invokeContext, cancel := context.WithDeadline(fn.Context(), deadline)
 	defer cancel()
 
 	lc := &lambdacontext.LambdaContext{

lambda/function_test.go

@@ -58,6 +58,32 @@ func TestInvoke(t *testing.T) {
 	assert.Equal(t, deadline.UnixNano(), responseValue)
 }
 
+func TestInvokeWithContext(t *testing.T) {
+	key := struct{}{}
+	srv := &Function{
+		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("!?!?!?!?!")
+			}),
+		context: context.WithValue(context.Background(), key, "dummy"),
+	}
+	deadline := time.Now()
+	var response messages.InvokeResponse
+	err := srv.Invoke(&messages.InvokeRequest{
+		Deadline: messages.InvokeRequest_Timestamp{
+			Seconds: deadline.Unix(),
+			Nanos:   int64(deadline.Nanosecond()),
+		}}, &response)
+	assert.NoError(t, err)
+	var responseValue int64
+	assert.NoError(t, json.Unmarshal(response.Payload, &responseValue))
+	assert.Equal(t, deadline.UnixNano(), responseValue)
+}
+
 type CustomError struct{}
 
 func (e CustomError) Error() string { return fmt.Sprintf("Something bad happened!") }