Add support for deadlines

Adds a generic `Metadata` field where we store the deadline for the
request, which is taken from the passed in context.
This deadline is then added to the context passed into the server
handler function.

This is pretty simnilar to how GRPC implements this, except GRPC doesn't
support deadlines, instead a timeout is specified where the timeouty is
the duration left until the deadline just before the client sends the
request... I chose to use a deadline since it seems like this is a more
precise implementation of a caller's intention.

Signed-off-by: Brian Goff <[email protected]>

Author: cpuguy83

client.go

@@ -21,6 +21,7 @@ import (
 	"io"
 	"net"
 	"os"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -78,14 +79,19 @@ func (c *Client) Call(ctx context.Context, service, method string, req, resp int
 
 	var (
 		creq = &Request{
-			Service: service,
-			Method:  method,
-			Payload: payload,
+			Service:  service,
+			Method:   method,
+			Payload:  payload,
+			Metadata: make(map[string]*Values),
 		}
 
 		cresp = &Response{}
 	)
 
+	if dl, ok := ctx.Deadline(); ok {
+		creq.Metadata[metadataDeadlineKey] = &Values{Values: []string{strconv.FormatInt(dl.UnixNano(), 10)}}
+	}
+
 	if err := c.dispatch(ctx, creq, cresp); err != nil {
 		return err
 	}
@@ -104,6 +110,7 @@ func (c *Client) Call(ctx context.Context, service, method string, req, resp int
 func (c *Client) dispatch(ctx context.Context, req *Request, resp *Response) error {
 	errs := make(chan error, 1)
 	call := &callRequest{
+		ctx:  context.TODO(),
 		req:  req,
 		resp: resp,
 		errs: errs,

server.go

@@ -21,16 +21,20 @@ import (
 	"io"
 	"math/rand"
 	"net"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
+	spb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
+const metadataDeadlineKey = "ttrpc-deadline"
+
 var (
 	ErrServerClosed = errors.New("ttrpc: server closed")
 )
@@ -414,6 +418,26 @@ func (c *serverConn) run(sctx context.Context) {
 		case request := <-requests:
 			active++
 			go func(id uint32) {
+				ctx, cancel, err := getRequestContext(ctx, request.req)
+				if err != nil {
+					resp := response{
+						id: id,
+						resp: &Response{
+							Status: &spb.Status{
+								Code:    int32(codes.InvalidArgument),
+								Message: err.Error(),
+							},
+						},
+					}
+					select {
+					case responses <- resp:
+					case <-done:
+					}
+					return
+				}
+
+				defer cancel()
+
 				p, status := c.server.services.call(ctx, request.req.Service, request.req.Method, request.req.Payload)
 				resp := &Response{
 					Status:  status.Proto(),
@@ -454,3 +478,27 @@ func (c *serverConn) run(sctx context.Context) {
 		}
 	}
 }
+
+var noopFunc = func() {}
+
+func getRequestContext(ctx context.Context, req *Request) (retCtx context.Context, cancel func(), err error) {
+	cancel = noopFunc
+	md := req.GetMetadata()
+
+	deadlineValues, ok := md[metadataDeadlineKey]
+	if !ok {
+		return ctx, cancel, nil
+	}
+
+	if len(deadlineValues) == 0 {
+		return ctx, cancel, nil
+	}
+
+	nano, err := strconv.ParseInt(deadlineValues[0], 10, 64)
+	if err != nil {
+		return ctx, cancel, errors.New("invalid value for deadline")
+	}
+
+	ctx, cancel = context.WithDeadline(ctx, time.Unix(0, nano))
+	return ctx, cancel, nil
+}

server_test.go

@@ -24,6 +24,7 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
@@ -57,7 +58,8 @@ func (tc *testingClient) Test(ctx context.Context, req *testPayload) (*testPaylo
 }
 
 type testPayload struct {
-	Foo string `protobuf:"bytes,1,opt,name=foo,proto3"`
+	Foo      string `protobuf:"bytes,1,opt,name=foo,proto3"`
+	Deadline int64  `protobuf:"varint,2,opt,name=deadline,proto3"`
 }
 
 func (r *testPayload) Reset()         { *r = testPayload{} }
@@ -68,7 +70,11 @@ func (r *testPayload) ProtoMessage()  {}
 type testingServer struct{}
 
 func (s *testingServer) Test(ctx context.Context, req *testPayload) (*testPayload, error) {
-	return &testPayload{Foo: strings.Repeat(req.Foo, 2)}, nil
+	tp := &testPayload{Foo: strings.Repeat(req.Foo, 2)}
+	if dl, ok := ctx.Deadline(); ok {
+		tp.Deadline = dl.UnixNano()
+	}
+	return tp, nil
 }
 
 // registerTestingService mocks more of what is generated code. Unlike grpc, we
@@ -376,6 +382,34 @@ func TestUnixSocketHandshake(t *testing.T) {
 	}
 }
 
+func TestServerRequestTimeout(t *testing.T) {
+	var (
+		ctx, cancel     = context.WithDeadline(context.Background(), time.Now().Add(10*time.Minute))
+		server          = mustServer(t)(NewServer())
+		addr, listener  = newTestListener(t)
+		testImpl        = &testingServer{}
+		client, cleanup = newTestClient(t, addr)
+		result          testPayload
+	)
+	defer cancel()
+	defer cleanup()
+	defer listener.Close()
+
+	registerTestingService(server, testImpl)
+
+	go server.Serve(ctx, listener)
+	defer server.Shutdown(ctx)
+
+	if err := client.Call(ctx, serviceName, "Test", &testPayload{}, &result); err != nil {
+		t.Fatalf("unexpected error making call: %v", err)
+	}
+
+	dl, _ := ctx.Deadline()
+	if result.Deadline != dl.UnixNano() {
+		t.Fatalf("expected deadline %v, actual: %v", dl, result.Deadline)
+	}
+}
+
 func BenchmarkRoundTrip(b *testing.B) {
 	var (
 		ctx             = context.Background()

types.go

@@ -23,15 +23,34 @@ import (
 )
 
 type Request struct {
-	Service string `protobuf:"bytes,1,opt,name=service,proto3"`
-	Method  string `protobuf:"bytes,2,opt,name=method,proto3"`
-	Payload []byte `protobuf:"bytes,3,opt,name=payload,proto3"`
+	Service  string             `protobuf:"bytes,1,opt,name=service,proto3"`
+	Method   string             `protobuf:"bytes,2,opt,name=method,proto3"`
+	Payload  []byte             `protobuf:"bytes,3,opt,name=payload,proto3"`
+	Metadata map[string]*Values `protobuf:"bytes,4,rep,name=metadata,proto3" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
 }
 
 func (r *Request) Reset()         { *r = Request{} }
 func (r *Request) String() string { return fmt.Sprintf("%+#v", r) }
 func (r *Request) ProtoMessage()  {}
 
+type Metadata map[string][]string
+
+func (r *Request) GetMetadata() Metadata {
+	md := make(Metadata)
+	for k, v := range r.Metadata {
+		md[k] = v.Values
+	}
+	return md
+}
+
+type Values struct {
+	Values []string `protobuf:"bytes,1,rep,name=values,proto3"`
+}
+
+func (v *Values) Reset()         { *v = Values{} }
+func (v *Values) String() string { return fmt.Sprintf("%+#v", v) }
+func (v *Values) ProtoMessage()  {}
+
 type Response struct {
 	Status  *spb.Status `protobuf:"bytes,1,opt,name=status,proto3"`
 	Payload []byte      `protobuf:"bytes,2,opt,name=payload,proto3"`