From 01ed7d87777efbbe3bc00045483b475699458c6e Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Mon, 2 Jul 2018 11:31:25 -0400
Subject: [PATCH] Add context to Serve

This is to seed a new context on the serve method

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 server.go      |  3 +--
 server_test.go | 22 ++++++++++++----------
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/server.go b/server.go
index 3797a09..1a518b1 100644
--- a/server.go
+++ b/server.go
@@ -67,12 +67,11 @@ func (s *Server) Register(name string, methods map[string]Method) {
 	s.services.register(name, methods)
 }
 
-func (s *Server) Serve(l net.Listener) error {
+func (s *Server) Serve(ctx context.Context, l net.Listener) error {
 	s.addListener(l)
 	defer s.closeListener(l)
 
 	var (
-		ctx        = context.Background()
 		backoff    time.Duration
 		handshaker = s.config.handshaker
 	)
diff --git a/server_test.go b/server_test.go
index 46b84a0..ec07f0e 100644
--- a/server_test.go
+++ b/server_test.go
@@ -107,7 +107,7 @@ func TestServer(t *testing.T) {
 
 	registerTestingService(server, testImpl)
 
-	go server.Serve(listener)
+	go server.Serve(ctx, listener)
 	defer server.Shutdown(ctx)
 
 	const calls = 2
@@ -134,7 +134,7 @@ func TestServerNotFound(t *testing.T) {
 	defer cleanup()
 	defer listener.Close()
 	go func() {
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	var tp testPayload
@@ -156,13 +156,14 @@ func TestServerNotFound(t *testing.T) {
 
 func TestServerListenerClosed(t *testing.T) {
 	var (
+		ctx         = context.Background()
 		server      = mustServer(t)(NewServer())
 		_, listener = newTestListener(t)
 		errs        = make(chan error, 1)
 	)
 
 	go func() {
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	if err := listener.Close(); err != nil {
@@ -211,7 +212,7 @@ func TestServerShutdown(t *testing.T) {
 	})
 
 	go func() {
-		serveErrs <- server.Serve(listener)
+		serveErrs <- server.Serve(ctx, listener)
 	}()
 
 	// send a series of requests that will get blocked
@@ -254,6 +255,7 @@ func TestServerShutdown(t *testing.T) {
 
 func TestServerClose(t *testing.T) {
 	var (
+		ctx         = context.Background()
 		server      = mustServer(t)(NewServer())
 		_, listener = newTestListener(t)
 		startClose  = make(chan struct{})
@@ -262,7 +264,7 @@ func TestServerClose(t *testing.T) {
 
 	go func() {
 		close(startClose)
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	<-startClose
@@ -289,7 +291,7 @@ func TestOversizeCall(t *testing.T) {
 	defer cleanup()
 	defer listener.Close()
 	go func() {
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	registerTestingService(server, &testingServer{})
@@ -324,7 +326,7 @@ func TestClientEOF(t *testing.T) {
 	defer cleanup()
 	defer listener.Close()
 	go func() {
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	registerTestingService(server, &testingServer{})
@@ -362,7 +364,7 @@ func TestUnixSocketHandshake(t *testing.T) {
 	defer cleanup()
 	defer listener.Close()
 	go func() {
-		errs <- server.Serve(listener)
+		errs <- server.Serve(ctx, listener)
 	}()
 
 	registerTestingService(server, &testingServer{})
@@ -389,7 +391,7 @@ func BenchmarkRoundTrip(b *testing.B) {
 
 	registerTestingService(server, testImpl)
 
-	go server.Serve(listener)
+	go server.Serve(ctx, listener)
 	defer server.Shutdown(ctx)
 
 	var tp testPayload
@@ -421,7 +423,7 @@ func BenchmarkRoundTripUnixSocketCreds(b *testing.B) {
 
 	registerTestingService(server, testImpl)
 
-	go server.Serve(listener)
+	go server.Serve(ctx, listener)
 	defer server.Shutdown(ctx)
 
 	var tp testPayload