Add support for graceful close timeouts via contect.Context

* `Close()` is now `Close(context.Context)` for sessions and links.
* `Client.Close` does not take a context as it does not currently
  block.
* Added `-count=1` to integration tests to prevent caching on Go1.10+.

Author: vcabbage

Makefile

@@ -20,7 +20,7 @@ test:
 	go test -tags gofuzz -v -race ./...
 
 integration:
-	go test -tags "integration pkgerrors" -v -race ./...
+	go test -tags "integration pkgerrors" -count=1 -v -race ./...
 
 coverage:
 	TEST_CORPUS=1 go test -tags "integration gofuzz" -cover -coverprofile=cover.out -v

README.md

@@ -36,6 +36,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"time"
 
 	"pack.ag/amqp"
 )
@@ -76,8 +77,8 @@ func main() {
 			log.Fatal("Sending message:", err)
 		}
 
+		sender.Close(ctx)
 		cancel()
-		sender.Close()
 	}
 
 	// Continuously read messages
@@ -90,9 +91,11 @@ func main() {
 		if err != nil {
 			log.Fatal("Creating receiver link:", err)
 		}
-
-		ctx, cancel := context.WithCancel(ctx)
-		defer cancel()
+		defer func() {
+			ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+			receiver.Close(ctx)
+			cancel()
+		}()
 
 		for {
 			// Receive next message

client.go

@@ -139,7 +139,7 @@ func (c *Client) NewSession() (*Session, error) {
 
 	begin, ok := fr.body.(*performBegin)
 	if !ok {
-		_ = s.Close() // deallocate session on error
+		_ = s.Close(context.Background()) // deallocate session on error
 		return nil, errorErrorf("unexpected begin response: %+v", fr.body)
 	}
 
@@ -193,10 +193,17 @@ func newSession(c *conn, channel uint16) *Session {
 	}
 }
 
-// Close closes the session.
-func (s *Session) Close() error {
+// Close gracefully closes the session.
+//
+// If ctx expires while waiting for servers response, ctx.Err() will be returned.
+// The session will continue to wait for the response until the Client is closed.
+func (s *Session) Close(ctx context.Context) error {
 	s.closeOnce.Do(func() { close(s.close) })
-	<-s.done
+	select {
+	case <-s.done:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 	if s.err == ErrSessionClosed {
 		return nil
 	}
@@ -351,10 +358,8 @@ func (s *Sender) Address() string {
 }
 
 // Close closes the Sender and AMQP link.
-func (s *Sender) Close() error {
-	// TODO: Should this timeout? Close() take a context? Use one of the
-	// other timeouts?
-	return s.link.Close()
+func (s *Sender) Close(ctx context.Context) error {
+	return s.link.Close(ctx)
 }
 
 // NewSender opens a new sender link on the session.
@@ -1045,9 +1050,17 @@ func (l *link) mux() {
 // close closes and requests deletion of the link.
 //
 // No operations on link are valid after close.
-func (l *link) Close() error {
+//
+// If ctx expires while waiting for servers response, ctx.Err() will be returned.
+// The session will continue to wait for the response until the Session or Client
+// is closed.
+func (l *link) Close(ctx context.Context) error {
 	l.closeOnce.Do(func() { close(l.close) })
-	<-l.done
+	select {
+	case <-l.done:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 	if l.err == ErrLinkClosed {
 		return nil
 	}
@@ -1350,7 +1363,7 @@ func (r *Receiver) Receive(ctx context.Context) (*Message, error) {
 		messageSize += len(fr.Payload)
 		if messageSize > maxMessageSize {
 			// TODO: send error
-			_ = r.Close()
+			_ = r.Close(ctx)
 			return nil, errorErrorf("received message larger than max size of ")
 		}
 
@@ -1389,10 +1402,12 @@ func (r *Receiver) Address() string {
 }
 
 // Close closes the Receiver and AMQP link.
-func (r *Receiver) Close() error {
-	// TODO: Should this timeout? Close() take a context? Use one of the
-	// other timeouts?
-	return r.link.Close()
+//
+// If ctx expires while waiting for servers response, ctx.Err() will be returned.
+// The session will continue to wait for the response until the Session or Client
+// is closed.
+func (r *Receiver) Close(ctx context.Context) error {
+	return r.link.Close(ctx)
 }
 
 type messageDisposition struct {

example_test.go

@@ -45,8 +45,8 @@ func Example() {
 			log.Fatal("Sending message:", err)
 		}
 
+		sender.Close(ctx)
 		cancel()
-		sender.Close()
 	}
 
 	// Continuously read messages
@@ -59,9 +59,11 @@ func Example() {
 		if err != nil {
 			log.Fatal("Creating receiver link:", err)
 		}
-
-		ctx, cancel := context.WithCancel(ctx)
-		defer cancel()
+		defer func() {
+			ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+			receiver.Close(ctx)
+			cancel()
+		}()
 
 		for {
 			// Receive next message

fuzz.go

@@ -37,9 +37,12 @@ func FuzzConn(data []byte) int {
 
 	msg.Accept()
 
-	// r.Close() // disabled until link close timeout implemented
+	ctx, close := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer close()
 
-	s.Close()
+	r.Close(ctx)
+
+	s.Close(ctx)
 
 	// Send
 	client, err = New(testconn.New(data),
@@ -66,9 +69,12 @@ func FuzzConn(data []byte) int {
 		return 0
 	}
 
-	// r.Close() // disabled until link close timeout implemented
+	ctx, close = context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer close()
+
+	r.Close(ctx)
 
-	s.Close()
+	s.Close(ctx)
 
 	return 1
 }