Merge pull request #119 from arangodb/bug-fix/stream-cursor-test

VST stream cursor test fixes & VST fail-quick fix

Author: ewoutp

Makefile

@@ -298,6 +298,7 @@ __test_go_test:
 		-e TEST_CONNECTION=$(TEST_CONNECTION) \
 		-e TEST_CVERSION=$(TEST_CVERSION) \
 		-e TEST_CONTENT_TYPE=$(TEST_CONTENT_TYPE) \
+		-e TEST_PPROF=$(TEST_PPROF) \
 		-w /usr/code/ \
 		golang:$(GOVERSION) \
 		go test $(TAGS) $(TESTOPTIONS) $(TESTVERBOSEOPTIONS) $(TESTS)

test/client_test.go

@@ -25,6 +25,7 @@ package test
 import (
 	"context"
 	"crypto/tls"
+	"log"
 	httplib "net/http"
 	"os"
 	"strconv"
@@ -33,14 +34,17 @@ import (
 	"testing"
 	"time"
 
+	_ "net/http/pprof"
+
 	driver "github.com/arangodb/go-driver"
 	"github.com/arangodb/go-driver/http"
 	"github.com/arangodb/go-driver/vst"
 	"github.com/arangodb/go-driver/vst/protocol"
 )
 
 var (
-	logEndpointsOnce sync.Once
+	logEndpointsOnce   sync.Once
+	runPProfServerOnce sync.Once
 )
 
 // skipBelowVersion skips the test if the current server version is less than
@@ -152,6 +156,16 @@ func createConnectionFromEnv(t testEnv) driver.Connection {
 
 // createClientFromEnv initializes a Client from information specified in environment variables.
 func createClientFromEnv(t testEnv, waitUntilReady bool, connection ...*driver.Connection) driver.Client {
+	runPProfServerOnce.Do(func() {
+		if os.Getenv("TEST_PPROF") != "" {
+			go func() {
+				// Start pprof server on port 6060
+				// To use it in the test, run a command like:
+				// docker exec -it go-driver-test sh -c "apk add -U curl && curl http://localhost:6060/debug/pprof/goroutine?debug=1"
+				log.Println(httplib.ListenAndServe("localhost:6060", nil))
+			}()
+		}
+	})
 	conn := createConnectionFromEnv(t)
 	if len(connection) == 1 {
 		*connection[0] = conn

test/cursor_test.go

@@ -261,6 +261,7 @@ func TestCreateStreamCursor(t *testing.T) {
 			t.Fatalf("Expected success, got %s", describe(err))
 		}
 	}
+	t.Log("Completed inserting 10k docs")
 
 	const expectedResults int = 10 * 10000
 	query := "FOR doc IN cursor_stream_test RETURN doc"
@@ -270,14 +271,10 @@ func TestCreateStreamCursor(t *testing.T) {
 	// create a bunch of read-only cursors
 	for i := 0; i < 10; i++ {
 		cursor, err := db.Query(ctx2, query, nil)
-		if err == nil {
-			// Close upon exit of the function
-			defer cursor.Close()
-		}
 		if err != nil {
-			t.Errorf("Expected success in query %d (%s), got '%s'", i, query, describe(err))
-			continue
+			t.Fatalf("Expected success in query %d (%s), got '%s'", i, query, describe(err))
 		}
+		defer cursor.Close()
 		count := cursor.Count()
 		if count != 0 {
 			t.Errorf("Expected count of 0, got %d in query %d (%s)", count, i, query)
@@ -294,25 +291,24 @@ func TestCreateStreamCursor(t *testing.T) {
 		cursors = append(cursors, cursor)
 	}
 
+	t.Logf("Created %d cursors", len(cursors))
+
 	// start a write query on the same collection inbetween
 	// contrary to normal cursors which are executed right
 	// away this will block until all read cursors are resolved
 	testReady := make(chan bool)
 	go func() {
 		query = "FOR doc IN 1..5 LET y = SLEEP(0.01) INSERT {name:'Peter', age:0} INTO cursor_stream_test"
 		cursor, err := db.Query(ctx2, query, nil) // should not return immediately
-		if err == nil {
-			// Close upon exit of the function
-			defer cursor.Close()
-		}
 		if err != nil {
-			t.Errorf("Expected success in write-query %s, got '%s'", query, describe(err))
+			t.Fatalf("Expected success in write-query %s, got '%s'", query, describe(err))
 		}
+		defer cursor.Close()
 
 		for cursor.HasMore() {
 			var data interface{}
 			if _, err := cursor.ReadDocument(ctx2, &data); err != nil {
-				t.Errorf("Failed to read document, err: %s", describe(err))
+				t.Fatalf("Failed to read document, err: %s", describe(err))
 			}
 		}
 		testReady <- true // signal write done
@@ -325,7 +321,7 @@ func TestCreateStreamCursor(t *testing.T) {
 			for cursor.HasMore() {
 				var user UserDoc
 				if _, err := cursor.ReadDocument(ctx2, &user); err != nil {
-					t.Errorf("Failed to result document %d: %s", i, describe(err))
+					t.Fatalf("Failed to result document %d: %s", i, describe(err))
 				}
 				readCount++
 			}
@@ -335,10 +331,9 @@ func TestCreateStreamCursor(t *testing.T) {
 
 	writeDone := false
 	readDone := false
-	deadline := time.Now().Add(time.Second * 30)
 	for {
 		select {
-		case <-time.After(time.Until(deadline)):
+		case <-ctx.Done():
 			t.Fatal("Timeout")
 		case v := <-testReady:
 			if v {

vst/connection.go

@@ -142,6 +142,9 @@ func (c *vstConnection) Do(ctx context.Context, req driver.Request) (driver.Resp
 
 // Do performs a given request, returning its response.
 func (c *vstConnection) do(ctx context.Context, req driver.Request, transport messageTransport) (driver.Response, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	vstReq, ok := req.(*vstRequest)
 	if !ok {
 		return nil, driver.WithStack(driver.InvalidArgumentError{Message: "request is not a *vstRequest"})
@@ -158,10 +161,16 @@ func (c *vstConnection) do(ctx context.Context, req driver.Request, transport me
 	vstReq.WroteRequest()
 
 	// Wait for response
-	msg, ok := <-resp
-	if !ok {
-		// Message was cancelled / timeout
-		return nil, driver.WithStack(context.DeadlineExceeded)
+	var msg protocol.Message
+	select {
+	case msg, ok = <-resp:
+		if !ok {
+			// Message was canceled / timeout
+			return nil, driver.WithStack(context.DeadlineExceeded)
+		}
+	case <-ctx.Done():
+		// Context canceled while waiting here
+		return nil, driver.WithStack(ctx.Err())
 	}
 
 	//fmt.Printf("Received msg: %d\n", msg.ID)