Merge pull request #3 from beyondnetworks/add-request-timeout

Add request timeout

Author: AdallomRoy

client.go

@@ -22,6 +22,7 @@ package grammes
 
 import (
 	"sync"
+	"time"
 
 	"github.com/northwesternmutual/grammes/gremconnect"
 	"github.com/northwesternmutual/grammes/gremerror"
@@ -31,6 +32,7 @@ import (
 
 // maxConCurrentMessages determines the size of the request channel.
 const maxConCurrentMessages = 3
+const defaultTimeout = time.Minute
 
 // Client is used to handle the graph, schema, connection,
 // and basic debug logging when querying the graph database.
@@ -60,6 +62,8 @@ type Client struct {
 	broken bool
 	// logger is used to log out debug statements and errors from the client.
 	logger logging.Logger
+	// requestTimeout is used for timeouting requests that a response is not received for
+	requestTimeout time.Duration
 }
 
 // setupClient default values some fields in the client.
@@ -71,6 +75,7 @@ func setupClient() *Client {
 		resultMessenger: &sync.Map{},
 		logger:          logging.NewNilLogger(),
 		gremlinVersion:  "3",
+		requestTimeout:  defaultTimeout,
 	}
 }
 

configuration.go

@@ -102,3 +102,11 @@ func WithReadingWait(interval time.Duration) ClientConfiguration {
 		c.conn.SetReadingWait(interval)
 	}
 }
+
+// WithRequestTimeout sets the timeout when
+// reading a request from the gremlin server
+func WithRequestTimeout(interval time.Duration) ClientConfiguration {
+	return func(c *Client) {
+		c.requestTimeout = interval
+	}
+}

response.go

@@ -22,6 +22,8 @@ package grammes
 
 import (
 	"encoding/json"
+	"errors"
+	"time"
 
 	"github.com/northwesternmutual/grammes/gremconnect"
 )
@@ -71,20 +73,32 @@ func (c *Client) retrieveResponse(id string) ([][]byte, error) {
 		dataPart    []byte
 	)
 
-	if n := <-notifier.(chan int); n == 1 {
-		if dataI, ok := c.results.Load(id); ok {
-			for _, d := range dataI.([]interface{}) {
-				if err, ok = d.(error); ok {
-					break
-				}
-				if dataPart, err = jsonMarshalData(d); err != nil {
-					break
+	timeout := make(chan bool, 1)
+
+	time.AfterFunc(c.requestTimeout, func() {
+		timeout <- true
+	})
+
+	for n := 1;n == 1; {
+		select {
+		case n = <-notifier.(chan int):
+
+			if dataI, ok := c.results.Load(id); ok {
+				for _, d := range dataI.([]interface{}) {
+					if err, ok = d.(error); ok {
+						break
+					}
+					if dataPart, err = jsonMarshalData(d); err != nil {
+						break
+					}
+					data = append(data, dataPart)
 				}
-				data = append(data, dataPart)
+				close(notifier.(chan int))
+				c.resultMessenger.Delete(id)
+				c.deleteResponse(id)
 			}
-			close(notifier.(chan int))
-			c.resultMessenger.Delete(id)
-			c.deleteResponse(id)
+		case <-timeout:
+			return nil, errors.New("request failed with timeout")
 		}
 	}