Implement graceful shutdown of agent.DataAccess (#1008)

* Implements graceful shutdown of agent.DataAccess and underlying Datastore/Logstore/MessageQueue

* adds tests for closing agent.DataAccess and Datastore

api/agent/agent.go

@@ -84,6 +84,7 @@ type Agent interface {
 	Submit(Call) error
 
 	// Close will wait for any outstanding calls to complete and then exit.
+	// Closing the agent will invoke Close on the underlying DataAccess.
 	// Close is not safe to be called from multiple threads.
 	io.Closer
 
@@ -209,6 +210,12 @@ func (a *agent) Close() error {
 		}
 	})
 
+	// shutdown any db/queue resources
+	// associated with DataAccess
+	daErr := a.da.Close()
+	if daErr != nil {
+		return daErr
+	}
 	return err
 }
 

api/agent/agent_test.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -61,6 +62,12 @@ func checkExpectedHeaders(t *testing.T, expectedHeaders http.Header, receivedHea
 	}
 }
 
+func checkClose(t *testing.T, a Agent) {
+	if err := a.Close(); err != nil {
+		t.Fatalf("Failed to close agent: %v", err)
+	}
+}
+
 func TestCallConfigurationRequest(t *testing.T) {
 	appName := "myapp"
 	path := "/"
@@ -94,7 +101,7 @@ func TestCallConfigurationRequest(t *testing.T) {
 	)
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	w := httptest.NewRecorder()
 
@@ -237,7 +244,7 @@ func TestCallConfigurationModel(t *testing.T) {
 	ds := datastore.NewMockInit()
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	callI, err := a.GetCall(FromModel(cm))
 	if err != nil {
@@ -308,7 +315,7 @@ func TestAsyncCallHeaders(t *testing.T) {
 	ds := datastore.NewMockInit()
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	callI, err := a.GetCall(FromModel(cm))
 	if err != nil {
@@ -434,7 +441,7 @@ func TestReqTooLarge(t *testing.T) {
 	cfg.MaxRequestSize = 5
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)), WithConfig(cfg))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	_, err = a.GetCall(FromModel(cm))
 	if err != models.ErrRequestContentTooBig {
@@ -487,7 +494,7 @@ func TestSubmitError(t *testing.T) {
 	ds := datastore.NewMockInit()
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -554,7 +561,7 @@ func TestHTTPWithoutContentLengthWorks(t *testing.T) {
 	)
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	bodOne := `{"echoContent":"yodawg"}`
 
@@ -617,7 +624,7 @@ func TestGetCallReturnsResourceImpossibility(t *testing.T) {
 	ds := datastore.NewMockInit()
 
 	a := New(NewCachedDataAccess(NewDirectDataAccess(ds, ds, new(mqs.Mock))))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	_, err := a.GetCall(FromModel(call))
 	if err != models.ErrCallTimeoutServerBusy {
@@ -723,7 +730,7 @@ func TestPipesAreClear(t *testing.T) {
 	)
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	// test read this body after 5s (after call times out) and make sure we don't get yodawg
 	// TODO could read after 10 seconds, to make sure the 2nd task's input stream isn't blocked
@@ -873,7 +880,7 @@ func TestPipesDontMakeSpuriousCalls(t *testing.T) {
 	)
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	bodOne := `{"echoContent":"yodawg"}`
 	req, err := http.NewRequest("GET", call.URL, strings.NewReader(bodOne))
@@ -979,7 +986,7 @@ func TestNBIOResourceTracker(t *testing.T) {
 	cfg.HotPoll = 20 * time.Millisecond
 
 	a := New(NewDirectDataAccess(ds, ds, new(mqs.Mock)), WithConfig(cfg))
-	defer a.Close()
+	defer checkClose(t, a)
 
 	reqCount := 20
 	errors := make(chan error, reqCount)
@@ -1029,3 +1036,43 @@ func TestNBIOResourceTracker(t *testing.T) {
 		t.Fatalf("Expected successes, but got %d", ok)
 	}
 }
+
+type closingDataAccess struct {
+	DataAccess
+	closeReturn error
+	closed      chan struct{}
+}
+
+func newClosingDataAccess(closeReturn error) *closingDataAccess {
+	ds := datastore.NewMockInit()
+	return &closingDataAccess{
+		DataAccess:  NewDirectDataAccess(ds, ds, new(mqs.Mock)),
+		closed:      make(chan struct{}),
+		closeReturn: closeReturn,
+	}
+
+}
+
+func (da *closingDataAccess) Close() error {
+	close(da.closed)
+	return da.closeReturn
+}
+
+func TestClosesDataAccess(t *testing.T) {
+	da := newClosingDataAccess(nil)
+
+	a := New(da)
+	checkClose(t, a)
+	<-da.closed
+}
+
+func TestCloseReturnsDataAccessError(t *testing.T) {
+	err := errors.New("foo")
+	da := newClosingDataAccess(err)
+	a := New(da)
+
+	if cerr := a.Close(); cerr != err {
+		t.Fatalf("Wrong error returned, expected %v but got %v", err, cerr)
+	}
+	<-da.closed
+}

api/agent/data_access.go

@@ -38,6 +38,11 @@ type DataAccess interface {
 	// Finish will notify the system that the Call has been processed, and
 	// fulfill the reservation in the queue if the call came from a queue.
 	Finish(ctx context.Context, mCall *models.Call, stderr io.Reader, async bool) error
+
+	// Close will wait for any pending operations to complete and
+	// shuts down connections to the underlying datastore/queue resources.
+	// Close is not safe to be called from multiple threads.
+	io.Closer
 }
 
 // CachedDataAccess wraps a DataAccess and caches the results of GetApp and GetRoute.
@@ -108,6 +113,11 @@ func (da *CachedDataAccess) GetRoute(ctx context.Context, appID string, routePat
 	return r.(*models.Route), nil
 }
 
+// Close invokes close on the underlying DataAccess
+func (da *CachedDataAccess) Close() error {
+	return da.DataAccess.Close()
+}
+
 type directDataAccess struct {
 	mq models.MessageQueue
 	ds models.Datastore
@@ -177,3 +187,18 @@ func (da *directDataAccess) Finish(ctx context.Context, mCall *models.Call, stde
 	}
 	return nil
 }
+
+// Close calls close on the underlying Datastore and MessageQueue. If the Logstore
+// and Datastore are different, it will call Close on the Logstore as well.
+func (da *directDataAccess) Close() error {
+	err := da.ds.Close()
+	if da.ds != da.ls {
+		if daErr := da.ls.Close(); daErr != nil {
+			err = daErr
+		}
+	}
+	if mqErr := da.mq.Close(); mqErr != nil {
+		err = mqErr
+	}
+	return err
+}

api/agent/hybrid/client.go

@@ -248,3 +248,7 @@ func (cl *client) once(ctx context.Context, request, result interface{}, method
 func (cl *client) url(args ...string) string {
 	return cl.base + strings.Join(args, "/")
 }
+
+func (cl *client) Close() error {
+	return nil
+}

api/agent/hybrid/nop.go

@@ -58,3 +58,7 @@ func (cl *nopDataStore) GetRoute(ctx context.Context, appName, route string) (*m
 	defer span.End()
 	return nil, errors.New("Should not call GetRoute on a NOP data store")
 }
+
+func (cl *nopDataStore) Close() error {
+	return nil
+}