Merge pull request #15 from digitalocean/mdl-context

ovsdb: implement cancelation via context

Author: mdlayher

ovsdb/client.go

@@ -15,7 +15,7 @@
 package ovsdb
 
 import (
-	"encoding/json"
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -42,11 +42,6 @@ type Client struct {
 	wg *sync.WaitGroup
 }
 
-type rpcResponse struct {
-	Result json.RawMessage
-	Error  error
-}
-
 // An OptionFunc is a function which can configure a Client.
 type OptionFunc func(c *Client) error
 
@@ -115,16 +110,21 @@ func (c *Client) Close() error {
 }
 
 // rpc performs a single RPC request, and checks the response for errors.
-func (c *Client) rpc(method string, out interface{}, args ...interface{}) error {
+func (c *Client) rpc(ctx context.Context, method string, out interface{}, args ...interface{}) error {
+	// Was the context canceled before sending the RPC?
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
 	// Unmarshal results into empty struct if no out specified.
 	if out == nil {
 		out = &struct{}{}
 	}
 
 	// Captures any OVSDB errors.
-	r := result{
-		Reply: out,
-	}
+	r := result{Reply: out}
 
 	req := jsonrpc.Request{
 		Method: method,
@@ -133,29 +133,23 @@ func (c *Client) rpc(method string, out interface{}, args ...interface{}) error
 	}
 
 	// Add callback for this RPC ID to return results via channel.
-	ch := make(chan rpcResponse, 0)
+	ch := make(chan rpcResponse, 1)
+	defer close(ch)
 	c.addCallback(req.ID, ch)
 
 	if err := c.c.Send(req); err != nil {
 		return err
 	}
 
-	// Wait for callback to fire.
-	res := <-ch
-	if err := res.Error; err != nil {
-		return err
+	// Await RPC completion or cancelation.
+	select {
+	case <-ctx.Done():
+		// RPC canceled; clean up callback.
+		return c.cancelCallback(ctx, req.ID)
+	case res := <-ch:
+		// RPC complete.
+		return rpcResult(res, &r)
 	}
-
-	if err := json.Unmarshal(res.Result, &r); err != nil {
-		return err
-	}
-
-	// OVSDB server returned an error, return it.
-	if r.Err != nil {
-		return r.Err
-	}
-
-	return nil
 }
 
 // listen starts an RPC receive loop that can return RPC results to
@@ -216,12 +210,23 @@ func (c *Client) doCallback(id int, res rpcResponse) {
 		return
 	}
 
-	// Return result, clean up channel, and remove this callback.
+	// Return result and remove this callback.
 	ch <- res
-	close(ch)
 	delete(c.callbacks, id)
 }
 
+// cancelCallback is invoked when an RPC is canceled by its context.
+func (c *Client) cancelCallback(ctx context.Context, id int) error {
+	// RPC canceled; acquire the callback mutex and clean up the callback
+	// for this RPC.
+	c.cbMu.Lock()
+	defer c.cbMu.Unlock()
+
+	delete(c.callbacks, id)
+
+	return ctx.Err()
+}
+
 func panicf(format string, a ...interface{}) {
 	panic(fmt.Sprintf(format, a...))
 }

ovsdb/client_integration_test.go

@@ -15,6 +15,7 @@
 package ovsdb_test
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"testing"
@@ -54,7 +55,7 @@ func TestClientIntegrationConcurrent(t *testing.T) {
 			<-sigC
 
 			for j := 0; j < 4; j++ {
-				_, err := c.ListDatabases()
+				_, err := c.ListDatabases(context.Background())
 				if err != nil {
 					panic(fmt.Sprintf("failed to query concurrently: %v", err))
 				}
@@ -72,7 +73,7 @@ func TestClientIntegrationConcurrent(t *testing.T) {
 }
 
 func testClientDatabases(t *testing.T, c *ovsdb.Client) {
-	dbs, err := c.ListDatabases()
+	dbs, err := c.ListDatabases(context.Background())
 	if err != nil {
 		t.Fatalf("failed to list databases: %v", err)
 	}

ovsdb/client_test.go

@@ -15,11 +15,13 @@
 package ovsdb_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"log"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/digitalocean/go-openvswitch/ovsdb"
 	"github.com/digitalocean/go-openvswitch/ovsdb/internal/jsonrpc"
@@ -37,7 +39,7 @@ func TestClientJSONRPCError(t *testing.T) {
 	})
 	defer done()
 
-	_, err := c.ListDatabases()
+	_, err := c.ListDatabases(context.Background())
 	if err == nil {
 		t.Fatal("expected an error, but none occurred")
 	}
@@ -58,7 +60,7 @@ func TestClientOVSDBError(t *testing.T) {
 	})
 	defer done()
 
-	_, err := c.ListDatabases()
+	_, err := c.ListDatabases(context.Background())
 	if err == nil {
 		t.Fatal("expected an error, but none occurred")
 	}
@@ -88,11 +90,60 @@ func TestClientBadCallback(t *testing.T) {
 		ID:     intPtr(10),
 	}
 
-	if _, err := c.ListDatabases(); err != nil {
+	if _, err := c.ListDatabases(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
 
+func TestClientContextCancelBeforeRPC(t *testing.T) {
+	// Context canceled before RPC even begins.
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c, _, done := testClient(t, func(_ jsonrpc.Request) jsonrpc.Response {
+		return jsonrpc.Response{
+			ID:     intPtr(1),
+			Result: mustMarshalJSON(t, []string{"foo"}),
+		}
+	})
+	defer done()
+
+	_, err := c.ListDatabases(ctx)
+	if err != context.Canceled {
+		t.Fatalf("expected context canceled error: %v", err)
+	}
+}
+
+func TestClientContextCancelDuringRPC(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping during short test run")
+	}
+
+	// Context canceled during long RPC.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c, _, done := testClient(t, func(_ jsonrpc.Request) jsonrpc.Response {
+		// RPC canceled; RPC server still processing.
+		// TODO(mdlayher): try to do something smarter than sleeping in a test.
+		cancel()
+		<-ctx.Done()
+
+		time.Sleep(500 * time.Millisecond)
+
+		return jsonrpc.Response{
+			ID:     intPtr(1),
+			Result: mustMarshalJSON(t, []string{"foo"}),
+		}
+	})
+	defer done()
+
+	_, err := c.ListDatabases(ctx)
+	if err != context.Canceled {
+		t.Fatalf("expected context canceled error: %v", err)
+	}
+}
+
 func testClient(t *testing.T, fn jsonrpc.TestFunc) (*ovsdb.Client, chan<- *jsonrpc.Response, func()) {
 	t.Helper()
 

ovsdb/internal/jsonrpc/testconn.go

@@ -57,7 +57,7 @@ func TestNetConn(t *testing.T, fn TestFunc) (net.Conn, chan<- *Response, func())
 	}
 
 	var wg sync.WaitGroup
-	wg.Add(2)
+	wg.Add(1)
 
 	notifC := make(chan *Response, 16)
 
@@ -67,7 +67,7 @@ func TestNetConn(t *testing.T, fn TestFunc) (net.Conn, chan<- *Response, func())
 		// Accept a single connection.
 		c, err := l.Accept()
 		if err != nil {
-			if strings.Contains(err.Error(), "use of closed network") {
+			if isNetworkCloseError(err) {
 				return
 			}
 
@@ -76,14 +76,28 @@ func TestNetConn(t *testing.T, fn TestFunc) (net.Conn, chan<- *Response, func())
 		defer c.Close()
 
 		dec := json.NewDecoder(c)
+
+		var encMu sync.RWMutex
 		enc := json.NewEncoder(c)
 
 		// Push RPC notifications to the client.
+		var notifWG sync.WaitGroup
+		notifWG.Add(1)
+		defer notifWG.Wait()
+
 		go func() {
-			defer wg.Done()
+			defer notifWG.Done()
 
 			for n := range notifC {
-				if err := enc.Encode(n); err != nil {
+				encMu.Lock()
+				err := enc.Encode(n)
+				encMu.Unlock()
+
+				if err != nil {
+					if isNetworkCloseError(err) {
+						return
+					}
+
 					panicf("failed to encode notification: %v", err)
 				}
 			}
@@ -93,15 +107,24 @@ func TestNetConn(t *testing.T, fn TestFunc) (net.Conn, chan<- *Response, func())
 		for {
 			var req Request
 			if err := dec.Decode(&req); err != nil {
-				if err == io.EOF {
+				if isNetworkCloseError(err) {
 					return
 				}
 
-				panicf("failed to decode request: %v", err)
+				panicf("failed to decode request: %#v", err)
 			}
 
 			res := fn(req)
-			if err := enc.Encode(res); err != nil {
+
+			encMu.Lock()
+			err := enc.Encode(res)
+			encMu.Unlock()
+
+			if err != nil {
+				if isNetworkCloseError(err) {
+					return
+				}
+
 				panicf("failed to encode response: %v", err)
 			}
 		}
@@ -124,3 +147,9 @@ func TestNetConn(t *testing.T, fn TestFunc) (net.Conn, chan<- *Response, func())
 func panicf(format string, a ...interface{}) {
 	panic(fmt.Sprintf(format, a...))
 }
+
+func isNetworkCloseError(err error) bool {
+	return err == io.EOF ||
+		strings.Contains(err.Error(), "use of closed network") ||
+		strings.Contains(err.Error(), "connection reset by peer")
+}

ovsdb/result.go

@@ -26,6 +26,31 @@ type result struct {
 	Err   *Error
 }
 
+// An rpcResponse is a response used in RPC callbacks.
+type rpcResponse struct {
+	Result json.RawMessage
+	Error  error
+}
+
+// rpcResult handles any errors from an rpcResponse and unmarshals results into
+// a result.
+func rpcResult(res rpcResponse, r *result) error {
+	if err := res.Error; err != nil {
+		return err
+	}
+
+	if err := json.Unmarshal(res.Result, &r); err != nil {
+		return err
+	}
+
+	// OVSDB server returned an error, return it.
+	if r.Err != nil {
+		return r.Err
+	}
+
+	return nil
+}
+
 // errPrefix is a prefix that occurs if an error is present in a JSON-RPC response.
 var errPrefix = []byte(`{"error":`)
 

ovsdb/rpc.go

@@ -14,10 +14,12 @@
 
 package ovsdb
 
+import "context"
+
 // ListDatabases returns the name of all databases known to the OVSDB server.
-func (c *Client) ListDatabases() ([]string, error) {
+func (c *Client) ListDatabases(ctx context.Context) ([]string, error) {
 	var dbs []string
-	if err := c.rpc("list_dbs", &dbs); err != nil {
+	if err := c.rpc(ctx, "list_dbs", &dbs); err != nil {
 		return nil, err
 	}
 

ovsdb/rpc_test.go

@@ -15,6 +15,7 @@
 package ovsdb_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/digitalocean/go-openvswitch/ovsdb/internal/jsonrpc"
@@ -40,7 +41,7 @@ func TestClientListDatabases(t *testing.T) {
 	})
 	defer done()
 
-	dbs, err := c.ListDatabases()
+	dbs, err := c.ListDatabases(context.Background())
 	if err != nil {
 		t.Fatalf("failed to list databases: %v", err)
 	}