Merge pull request #81 from sideshow/feature-context

Feature request context

Author: sideshow

.travis.yml

@@ -15,6 +15,7 @@ before_install:
 
 install:
   - go get golang.org/x/net/http2
+  - go get golang.org/x/net/context
   - go get golang.org/x/crypto/pkcs12
 
 os:

README.md

@@ -124,6 +124,20 @@ if res.Sent() {
 }
 ```
 
+## Context & Timeouts
+
+For better control over request cancelations and timeouts APNS/2 supports
+contexts. Using a context can be helpful if you want to cancel all pushes when
+the parent process is cancelled, or need finer grained control over individual
+push timeouts. See the [Google post](https://blog.golang.org/context) for more
+information on contexts.
+
+```go
+ctx, cancel = context.WithTimeout(context.Background(), 10 * time.Second)
+res, err := client.PushWithContext(ctx, notification)
+defer cancel()
+```
+
 ## Command line tool
 
 APNS/2 has a command line tool that can be installed with `go get github.com/sideshow/apns2/apns2`. Usage:

client.go

@@ -32,7 +32,7 @@ var (
 	// HTTPClientTimeout specifies a time limit for requests made by the
 	// HTTPClient. The timeout includes connection time, any redirects,
 	// and reading the response body.
-	HTTPClientTimeout = 30 * time.Second
+	HTTPClientTimeout = 60 * time.Second
 )
 
 // Client represents a connection with the APNs
@@ -93,17 +93,32 @@ func (c *Client) Production() *Client {
 // transparently before sending the notification. It will return a Response
 // indicating whether the notification was accepted or rejected by the APNs
 // gateway, or an error if something goes wrong.
+//
+// Use PushWithContext if you need better cancelation and timeout control.
 func (c *Client) Push(n *Notification) (*Response, error) {
-	payload, err := json.Marshal(n)
+	return c.PushWithContext(nil, n)
+}
 
+// PushWithContext sends a Notification to the APNs gateway. Context carries a
+// deadline and a cancelation signal and allows you to close long running
+// requests when the context timeout is exceeded. Context can be nil, for
+// backwards compatibility.
+//
+// If the underlying http.Client is not currently connected, this method will
+// attempt to reconnect transparently before sending the notification. It will
+// return a Response indicating whether the notification was accepted or
+// rejected by the APNs gateway, or an error if something goes wrong.
+func (c *Client) PushWithContext(ctx Context, n *Notification) (*Response, error) {
+	payload, err := json.Marshal(n)
 	if err != nil {
 		return nil, err
 	}
 
 	url := fmt.Sprintf("%v/3/device/%v", c.Host, n.DeviceToken)
 	req, _ := http.NewRequest("POST", url, bytes.NewBuffer(payload))
 	setHeaders(req, n)
-	httpRes, err := c.HTTPClient.Do(req)
+
+	httpRes, err := c.requestWithContext(ctx, req)
 	if err != nil {
 		return nil, err
 	}

client_go16.go

@@ -0,0 +1,25 @@
+// +build go1.6,!go1.7
+
+package apns2
+
+import (
+	"net/http"
+
+	"golang.org/x/net/context"
+	"golang.org/x/net/context/ctxhttp"
+)
+
+// A Context carries a deadline, a cancelation signal, and other values across
+// API boundaries.
+//
+// Context's methods may be called by multiple goroutines simultaneously.
+type Context interface {
+	context.Context
+}
+
+func (c *Client) requestWithContext(ctx Context, req *http.Request) (*http.Response, error) {
+	if ctx != nil {
+		return ctxhttp.Do(ctx, c.HTTPClient, req)
+	}
+	return c.HTTPClient.Do(req)
+}

client_go16_test.go

@@ -0,0 +1,48 @@
+// +build go1.6
+// +build !go1.7
+
+package apns2_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/net/context"
+)
+
+func TestClientPushWithContextWithTimeout(t *testing.T) {
+	const timeout = time.Nanosecond
+	n := mockNotification()
+	var apnsID = "02ABC856-EF8D-4E49-8F15-7B8A61D978D6"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		w.Header().Set("apns-id", apnsID)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	time.Sleep(timeout)
+	res, err := mockClient(server.URL).PushWithContext(ctx, n)
+	assert.Error(t, err)
+	assert.Nil(t, res)
+	cancel()
+}
+
+func TestClientPushWithContext(t *testing.T) {
+	n := mockNotification()
+	var apnsID = "02ABC856-EF8D-4E49-8F15-7B8A61D978D6"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		w.Header().Set("apns-id", apnsID)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	res, err := mockClient(server.URL).PushWithContext(context.Background(), n)
+	assert.Nil(t, err)
+	assert.Equal(t, res.ApnsID, apnsID)
+}

client_go17.go

@@ -0,0 +1,23 @@
+// +build go1.7
+
+package apns2
+
+import (
+	"context"
+	"net/http"
+)
+
+// A Context carries a deadline, a cancelation signal, and other values across
+// API boundaries.
+//
+// Context's methods may be called by multiple goroutines simultaneously.
+type Context interface {
+	context.Context
+}
+
+func (c *Client) requestWithContext(ctx Context, req *http.Request) (*http.Response, error) {
+	if ctx != nil {
+		req = req.WithContext(ctx)
+	}
+	return c.HTTPClient.Do(req)
+}

client_go17_test.go

@@ -0,0 +1,47 @@
+// +build go1.7
+
+package apns2_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestClientPushWithContextWithTimeout(t *testing.T) {
+	const timeout = time.Nanosecond
+	n := mockNotification()
+	var apnsID = "02ABC856-EF8D-4E49-8F15-7B8A61D978D6"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		w.Header().Set("apns-id", apnsID)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	time.Sleep(timeout)
+	res, err := mockClient(server.URL).PushWithContext(ctx, n)
+	assert.Error(t, err)
+	assert.Nil(t, res)
+	cancel()
+}
+
+func TestClientPushWithContext(t *testing.T) {
+	n := mockNotification()
+	var apnsID = "02ABC856-EF8D-4E49-8F15-7B8A61D978D6"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		w.Header().Set("apns-id", apnsID)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	res, err := mockClient(server.URL).PushWithContext(context.Background(), n)
+	assert.Nil(t, err)
+	assert.Equal(t, res.ApnsID, apnsID)
+}

client_test.go

@@ -69,8 +69,8 @@ func TestClientBadTransportError(t *testing.T) {
 }
 
 func TestClientNameToCertificate(t *testing.T) {
-	certificate, _ := certificate.FromP12File("certificate/_fixtures/certificate-valid.p12", "")
-	client := apns.NewClient(certificate)
+	crt, _ := certificate.FromP12File("certificate/_fixtures/certificate-valid.p12", "")
+	client := apns.NewClient(crt)
 	name := client.HTTPClient.Transport.(*http2.Transport).TLSClientConfig.NameToCertificate
 	assert.Len(t, name, 1)
 
@@ -82,8 +82,8 @@ func TestClientNameToCertificate(t *testing.T) {
 
 func TestDialTLSTimeout(t *testing.T) {
 	apns.TLSDialTimeout = 1 * time.Millisecond
-	certificate, _ := certificate.FromP12File("certificate/_fixtures/certificate-valid.p12", "")
-	client := apns.NewClient(certificate)
+	crt, _ := certificate.FromP12File("certificate/_fixtures/certificate-valid.p12", "")
+	client := apns.NewClient(crt)
 	dialTLS := client.HTTPClient.Transport.(*http2.Transport).DialTLS
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {