Add Retryable Writes

GODRIVER-53

Change-Id: I3613b72b45c311997e399635180e4fec9705ae8e

Author: rfblue2

core/command/command.go

@@ -17,6 +17,108 @@ import (
 	"github.com/mongodb/mongo-go-driver/core/writeconcern"
 )
 
+// DecodeError attempts to decode the wiremessage as an error
+func DecodeError(wm wiremessage.WireMessage) error {
+	var rdr bson.Reader
+	switch msg := wm.(type) {
+	case wiremessage.Msg:
+		for _, section := range msg.Sections {
+			switch converted := section.(type) {
+			case wiremessage.SectionBody:
+				rdr = converted.Document
+			}
+		}
+	case wiremessage.Reply:
+		if msg.ResponseFlags&wiremessage.QueryFailure != wiremessage.QueryFailure {
+			return nil
+		}
+		rdr = msg.Documents[0]
+	}
+
+	_, err := rdr.Validate()
+	if err != nil {
+		return nil
+	}
+
+	extractedError := extractError(rdr)
+
+	// If parsed successfully return the error
+	if _, ok := extractedError.(Error); ok {
+		return err
+	}
+
+	return nil
+}
+
+// helper method to extract an error from a reader if there is one; first returned item is the
+// error if it exists, the second holds parsing errors
+func extractError(rdr bson.Reader) error {
+	var errmsg, codeName string
+	var code int32
+	var labels []string
+	itr, err := rdr.Iterator()
+	if err != nil {
+		return err
+	}
+
+	for itr.Next() {
+		elem := itr.Element()
+		switch elem.Key() {
+		case "ok":
+			switch elem.Value().Type() {
+			case bson.TypeInt32:
+				if elem.Value().Int32() == 1 {
+					return nil
+				}
+			case bson.TypeInt64:
+				if elem.Value().Int64() == 1 {
+					return nil
+				}
+			case bson.TypeDouble:
+				if elem.Value().Double() == 1 {
+					return nil
+				}
+			}
+		case "errmsg":
+			if str, okay := elem.Value().StringValueOK(); okay {
+				errmsg = str
+			}
+		case "codeName":
+			if str, okay := elem.Value().StringValueOK(); okay {
+				codeName = str
+			}
+		case "code":
+			if c, okay := elem.Value().Int32OK(); okay {
+				code = c
+			}
+		case "errorLabels":
+			if arr, okay := elem.Value().MutableArrayOK(); okay {
+				iter, err := arr.Iterator()
+				if err != nil {
+					continue
+				}
+				for iter.Next() {
+					if str, ok := iter.Value().StringValueOK(); ok {
+						labels = append(labels, str)
+					}
+				}
+
+			}
+		}
+	}
+
+	if errmsg == "" {
+		errmsg = "command failed"
+	}
+
+	return Error{
+		Code:    code,
+		Message: errmsg,
+		Name:    codeName,
+		Labels:  labels,
+	}
+}
+
 func responseClusterTime(response bson.Reader) *bson.Document {
 	clusterTime, err := response.Lookup("$clusterTime")
 	if err != nil {

core/command/errors.go

@@ -13,6 +13,7 @@ import (
 	"strings"
 
 	"github.com/mongodb/mongo-go-driver/bson"
+	"github.com/mongodb/mongo-go-driver/core/result"
 )
 
 var (
@@ -116,6 +117,20 @@ func (e Error) Retryable() bool {
 	return false
 }
 
+// IsWriteConcernErrorRetryable returns true if the write concern error is retryable.
+func IsWriteConcernErrorRetryable(wce *result.WriteConcernError) bool {
+	for _, code := range retryableCodes {
+		if int32(wce.Code) == code {
+			return true
+		}
+	}
+	if strings.Contains(wce.ErrMsg, "not master") || strings.Contains(wce.ErrMsg, "node is recovering") {
+		return true
+	}
+
+	return false
+}
+
 // IsNotFound indicates if the error is from a namespace not being found.
 func IsNotFound(err error) bool {
 	e, ok := err.(Error)

core/command/insert.go

@@ -36,6 +36,7 @@ type Insert struct {
 	WriteConcern *writeconcern.WriteConcern
 	Session      *session.Client
 
+	batches         []*Write
 	result          result.Insert
 	err             error
 	continueOnError bool
@@ -89,13 +90,13 @@ splitInserts:
 
 // Encode will encode this command into a wire message for the given server description.
 func (i *Insert) Encode(desc description.SelectedServer) ([]wiremessage.WireMessage, error) {
-	cmds, err := i.encode(desc)
+	err := i.encode(desc)
 	if err != nil {
 		return nil, err
 	}
 
-	wms := make([]wiremessage.WireMessage, len(cmds))
-	for _, cmd := range cmds {
+	wms := make([]wiremessage.WireMessage, len(i.batches))
+	for _, cmd := range i.batches {
 		wm, err := cmd.Encode(desc)
 		if err != nil {
 			return nil, err
@@ -143,23 +144,21 @@ func (i *Insert) encodeBatch(docs []*bson.Document, desc description.SelectedSer
 	}, nil
 }
 
-func (i *Insert) encode(desc description.SelectedServer) ([]*Write, error) {
-	out := []*Write{}
+func (i *Insert) encode(desc description.SelectedServer) error {
 	batches, err := i.split(int(desc.MaxBatchCount), int(desc.MaxDocumentSize))
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	for _, docs := range batches {
 		cmd, err := i.encodeBatch(docs, desc)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
-		out = append(out, cmd)
+		i.batches = append(i.batches, cmd)
 	}
-
-	return out, nil
+	return nil
 }
 
 // Decode will decode the wire message using the provided server description. Errors during decoding
@@ -193,14 +192,25 @@ func (i *Insert) Err() error { return i.err }
 // RoundTrip handles the execution of this command using the provided wiremessage.ReadWriter.
 func (i *Insert) RoundTrip(ctx context.Context, desc description.SelectedServer, rw wiremessage.ReadWriter) (result.Insert, error) {
 	res := result.Insert{}
-	cmds, err := i.encode(desc)
-	if err != nil {
-		return res, err
+	if i.batches == nil {
+		err := i.encode(desc)
+		if err != nil {
+			return res, err
+		}
 	}
 
-	for _, cmd := range cmds {
+	// hold onto txnNumber, reset it when loop exits to ensure reuse of same
+	// transaction number if retry is needed
+	var txnNumber int64
+	if i.Session != nil && i.Session.RetryWrite {
+		txnNumber = i.Session.TxnNumber
+	}
+	for j, cmd := range i.batches {
 		rdr, err := cmd.RoundTrip(ctx, desc, rw)
 		if err != nil {
+			if i.Session != nil && i.Session.RetryWrite {
+				i.Session.TxnNumber = txnNumber + int64(j)
+			}
 			return res, err
 		}
 
@@ -213,13 +223,29 @@ func (i *Insert) RoundTrip(ctx context.Context, desc description.SelectedServer,
 
 		if r.WriteConcernError != nil {
 			res.WriteConcernError = r.WriteConcernError
+			if i.Session != nil && i.Session.RetryWrite {
+				i.Session.TxnNumber = txnNumber
+				return res, nil // report writeconcernerror for retry
+			}
 		}
 
 		res.N += r.N
 
 		if !i.continueOnError && len(res.WriteErrors) > 0 {
 			return res, nil
 		}
+
+		// Increment txnNumber for each batch
+		if i.Session != nil && i.Session.RetryWrite {
+			i.Session.IncrementTxnNumber()
+			i.batches = i.batches[1:] // if batch encoded successfully, remove it from the slice
+		}
+	}
+
+	if i.Session != nil && i.Session.RetryWrite {
+		// if retryable write succeeded, transaction number will be incremented one extra time,
+		// so we decrement it here
+		i.Session.TxnNumber--
 	}
 
 	return res, nil

core/command/opmsg.go

@@ -42,73 +42,9 @@ func decodeCommandOpMsg(msg wiremessage.Msg) (bson.Reader, error) {
 		return nil, NewCommandResponseError("malformed OP_MSG: invalid document", err)
 	}
 
-	ok := false
-	var errmsg, codeName string
-	var code int32
-	var labels []string
-	itr, err := rdr.Iterator()
+	err = extractError(rdr)
 	if err != nil {
-		return nil, NewCommandResponseError("malformed OP_MSG: cannot iterate document", err)
-	}
-
-	for itr.Next() {
-		elem := itr.Element()
-		switch elem.Key() {
-		case "ok":
-			switch elem.Value().Type() {
-			case bson.TypeInt32:
-				if elem.Value().Int32() == 1 {
-					ok = true
-				}
-			case bson.TypeInt64:
-				if elem.Value().Int64() == 1 {
-					ok = true
-				}
-			case bson.TypeDouble:
-				if elem.Value().Double() == 1 {
-					ok = true
-				}
-			}
-		case "errmsg":
-			if str, okay := elem.Value().StringValueOK(); okay {
-				errmsg = str
-			}
-		case "codeName":
-			if str, okay := elem.Value().StringValueOK(); okay {
-				codeName = str
-			}
-		case "code":
-			if c, okay := elem.Value().Int32OK(); okay {
-				code = c
-			}
-		case "errorLabels":
-			if arr, okay := elem.Value().MutableArrayOK(); okay {
-				iter, err := arr.Iterator()
-				if err != nil {
-					continue
-				}
-				for iter.Next() {
-					if str, ok := iter.Value().StringValueOK(); ok {
-						labels = append(labels, str)
-					}
-				}
-
-			}
-		}
-	}
-
-	if !ok {
-		if errmsg == "" {
-			errmsg = "command failed"
-		}
-
-		return nil, Error{
-			Code:    code,
-			Message: errmsg,
-			Name:    codeName,
-			Labels:  labels,
-		}
+		return nil, err
 	}
-
 	return rdr, nil
 }

core/command/opreply.go

@@ -35,56 +35,9 @@ func decodeCommandOpReply(reply wiremessage.Reply) (bson.Reader, error) {
 		}
 	}
 
-	ok := false
-	var errmsg, codeName string
-	var code int32
-	itr, err := rdr.Iterator()
+	err = extractError(rdr)
 	if err != nil {
-		return nil, NewCommandResponseError("malformed OP_REPLY: cannot iterate document", err)
+		return nil, err
 	}
-	for itr.Next() {
-		elem := itr.Element()
-		switch elem.Key() {
-		case "ok":
-			switch elem.Value().Type() {
-			case bson.TypeInt32:
-				if elem.Value().Int32() == 1 {
-					ok = true
-				}
-			case bson.TypeInt64:
-				if elem.Value().Int64() == 1 {
-					ok = true
-				}
-			case bson.TypeDouble:
-				if elem.Value().Double() == 1 {
-					ok = true
-				}
-			}
-		case "errmsg":
-			if str, okay := elem.Value().StringValueOK(); okay {
-				errmsg = str
-			}
-		case "codeName":
-			if str, okay := elem.Value().StringValueOK(); okay {
-				codeName = str
-			}
-		case "code":
-			if c, okay := elem.Value().Int32OK(); okay {
-				code = c
-			}
-		}
-	}
-
-	if !ok {
-		if errmsg == "" {
-			errmsg = "command failed"
-		}
-		return nil, Error{
-			Code:    code,
-			Message: errmsg,
-			Name:    codeName,
-		}
-	}
-
 	return rdr, nil
 }