httpu: add context.Context and related interface

This adds a new interface for httpu that supports a Context, and uses
that context to set a deadline/timeout and also cancel the request if
the context is canceled. Additionally, add a new method to the SSDP
package that takes a ClientInterfaceCtx.

Updates #55

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>

httpu/httpu.go

@@ -3,6 +3,7 @@ package httpu
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"log"
@@ -26,6 +27,27 @@ type ClientInterface interface {
 	) ([]*http.Response, error)
 }
 
+// ClientInterfaceCtx is the equivalent of ClientInterface, except with methods
+// taking a context.Context parameter.
+type ClientInterfaceCtx interface {
+	// DoWithContext performs a request. If the input request has a
+	// deadline, then that value will be used as the timeout for how long
+	// to wait before returning the responses that were received. If the
+	// request's context is canceled, this method will return immediately.
+	//
+	// If the request's context is never canceled, and does not have a
+	// deadline, then this function WILL NEVER RETURN. You MUST set an
+	// appropriate deadline on the context, or otherwise cancel it when you
+	// want to finish an operation.
+	//
+	// An error is only returned for failing to send the request. Failures
+	// in receipt simply do not add to the resulting responses.
+	DoWithContext(
+		req *http.Request,
+		numSends int,
+	) ([]*http.Response, error)
+}
+
 // HTTPUClient is a client for dealing with HTTPU (HTTP over UDP). Its typical
 // function is for HTTPMU, and particularly SSDP.
 type HTTPUClient struct {
@@ -34,6 +56,7 @@ type HTTPUClient struct {
 }
 
 var _ ClientInterface = &HTTPUClient{}
+var _ ClientInterfaceCtx = &HTTPUClient{}
 
 // NewHTTPUClient creates a new HTTPUClient, opening up a new UDP socket for the
 // purpose.
@@ -75,6 +98,25 @@ func (httpu *HTTPUClient) Do(
 	req *http.Request,
 	timeout time.Duration,
 	numSends int,
+) ([]*http.Response, error) {
+	ctx := req.Context()
+	if timeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+		req = req.WithContext(ctx)
+	}
+
+	return httpu.DoWithContext(req, numSends)
+}
+
+// DoWithContext implements ClientInterfaceCtx.DoWithContext.
+//
+// Make sure to read the documentation on the ClientInterfaceCtx interface
+// regarding cancellation!
+func (httpu *HTTPUClient) DoWithContext(
+	req *http.Request,
+	numSends int,
 ) ([]*http.Response, error) {
 	httpu.connLock.Lock()
 	defer httpu.connLock.Unlock()
@@ -101,10 +143,28 @@ func (httpu *HTTPUClient) Do(
 	if err != nil {
 		return nil, err
 	}
-	if err = httpu.conn.SetDeadline(time.Now().Add(timeout)); err != nil {
-		return nil, err
+
+	// Handle context deadline/timeout
+	ctx := req.Context()
+	deadline, ok := ctx.Deadline()
+	if ok {
+		if err = httpu.conn.SetDeadline(deadline); err != nil {
+			return nil, err
+		}
 	}
 
+	// Handle context cancelation
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		select {
+		case <-ctx.Done():
+			// if context is cancelled, stop any connections by setting time in the past.
+			httpu.conn.SetDeadline(time.Now().Add(-time.Second))
+		case <-done:
+		}
+	}()
+
 	// Send request.
 	for i := 0; i < numSends; i++ {
 		if n, err := httpu.conn.WriteTo(requestBuf.Bytes(), destAddr); err != nil {

httpu/multiclient.go

@@ -49,14 +49,14 @@ func (mc *MultiClient) Do(
 }
 
 func (mc *MultiClient) sendRequests(
-	results chan<-[]*http.Response,
+	results chan<- []*http.Response,
 	req *http.Request,
 	timeout time.Duration,
 	numSends int,
 ) error {
 	tasks := &errgroup.Group{}
 	for _, d := range mc.delegates {
-		d := d  // copy for closure
+		d := d // copy for closure
 		tasks.Go(func() error {
 			responses, err := d.Do(req, timeout, numSends)
 			if err != nil {
@@ -68,3 +68,65 @@ func (mc *MultiClient) sendRequests(
 	}
 	return tasks.Wait()
 }
+
+// MultiClientCtx dispatches requests out to all the delegated clients.
+type MultiClientCtx struct {
+	// The HTTPU clients to delegate to.
+	delegates []ClientInterfaceCtx
+}
+
+var _ ClientInterfaceCtx = &MultiClientCtx{}
+
+// NewMultiClient creates a new MultiClient that delegates to all the given
+// clients.
+func NewMultiClientCtx(delegates []ClientInterfaceCtx) *MultiClientCtx {
+	return &MultiClientCtx{
+		delegates: delegates,
+	}
+}
+
+// DoWithContext implements ClientInterfaceCtx.DoWithContext.
+func (mc *MultiClientCtx) DoWithContext(
+	req *http.Request,
+	numSends int,
+) ([]*http.Response, error) {
+	tasks, ctx := errgroup.WithContext(req.Context())
+	req = req.WithContext(ctx) // so we cancel if the errgroup errors
+	results := make(chan []*http.Response)
+
+	// For each client, send the request to it and collect results.
+	tasks.Go(func() error {
+		defer close(results)
+		return mc.sendRequestsCtx(results, req, numSends)
+	})
+
+	var responses []*http.Response
+	tasks.Go(func() error {
+		for rs := range results {
+			responses = append(responses, rs...)
+		}
+		return nil
+	})
+
+	return responses, tasks.Wait()
+}
+
+func (mc *MultiClientCtx) sendRequestsCtx(
+	results chan<- []*http.Response,
+	req *http.Request,
+	numSends int,
+) error {
+	tasks := &errgroup.Group{}
+	for _, d := range mc.delegates {
+		d := d // copy for closure
+		tasks.Go(func() error {
+			responses, err := d.DoWithContext(req, numSends)
+			if err != nil {
+				return err
+			}
+			results <- responses
+			return nil
+		})
+	}
+	return tasks.Wait()
+}