Merge pull request #17 from digitalocean/mdl-client-stats

 ovsdb: add ClientStats, fix callback leak

Author: mdlayher

ovsdb/client.go

@@ -109,6 +109,28 @@ func (c *Client) Close() error {
 	return err
 }
 
+// Stats returns a ClientStats with current statistics for the Cient.
+func (c *Client) Stats() ClientStats {
+	var s ClientStats
+
+	c.cbMu.RLock()
+	defer c.cbMu.RUnlock()
+
+	s.Callbacks.Current = len(c.callbacks)
+
+	return s
+}
+
+// ClientStats contains statistics about a Client.
+type ClientStats struct {
+	// Statistics about the Client's internal callbacks.
+	Callbacks struct {
+		// The number of callback hooks currently registered and waiting
+		// for RPC responses.
+		Current int
+	}
+}
+
 // rpc performs a single RPC request, and checks the response for errors.
 func (c *Client) rpc(ctx context.Context, method string, out interface{}, args ...interface{}) error {
 	// Was the context canceled before sending the RPC?
@@ -139,14 +161,28 @@ func (c *Client) rpc(ctx context.Context, method string, out interface{}, args .
 		Response: ch,
 	})
 
+	// Ensure that the callback is always cleaned up on return from this function.
+	// Note that this will result in the callback being deleted twice if the RPC
+	// returns successfully, but that's okay; it's a no-op.
+	//
+	// TODO(mdlayher): a more robust solution around callback map modifications.
+	defer func() {
+		c.cbMu.Lock()
+		defer c.cbMu.Unlock()
+
+		delete(c.callbacks, req.ID)
+	}()
+
 	if err := c.c.Send(req); err != nil {
 		return err
 	}
 
 	// Await RPC completion or cancelation.
 	select {
 	case <-ctx.Done():
-		// RPC canceled.  Producer cleans up the callback.
+		// RPC canceled.  The callback is cleaned up by deferred function in
+		// case no message ever arrives with its request ID, so we don't leak
+		// callbacks.
 		return ctx.Err()
 	case res, ok := <-ch:
 		if !ok {

ovsdb/client_test.go

@@ -144,6 +144,44 @@ func TestClientContextCancelDuringRPC(t *testing.T) {
 	}
 }
 
+func TestClientLeakCallbacks(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping during short test run")
+	}
+
+	c, _, done := testClient(t, func(_ jsonrpc.Request) jsonrpc.Response {
+		// Only respond with messages that don't match an incoming request.
+		return jsonrpc.Response{
+			ID:     intPtr(100),
+			Result: mustMarshalJSON(t, []string{"foo"}),
+		}
+	})
+	defer done()
+
+	// Expect no callbacks registered before RPCs, and none after RPCs time out.
+	var want ovsdb.ClientStats
+	want.Callbacks.Current = 0
+
+	if diff := cmp.Diff(want, c.Stats()); diff != "" {
+		t.Fatalf("unexpected starting client stats (-want +got):\n%s", diff)
+	}
+
+	for i := 0; i < 5; i++ {
+		// Give enough time for an RPC to be sent so we don't immediately time out.
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		_, err := c.ListDatabases(ctx)
+		if err != context.DeadlineExceeded {
+			t.Fatalf("expected context deadline exceeded error: %v", err)
+		}
+	}
+
+	if diff := cmp.Diff(want, c.Stats()); diff != "" {
+		t.Fatalf("unexpected ending client stats (-want +got):\n%s", diff)
+	}
+}
+
 func testClient(t *testing.T, fn jsonrpc.TestFunc) (*ovsdb.Client, chan<- *jsonrpc.Response, func()) {
 	t.Helper()