GODRIVER-124: batch splitting implementation

Change-Id: I4371a92c1cd74769dd25ffac272da5ca018afa1d

Author: tychoish

core/command/insert.go

@@ -16,6 +16,10 @@ import (
 	"github.com/mongodb/mongo-go-driver/core/wiremessage"
 )
 
+// this is the amount of reserved buffer space in a message that the
+// driver reserves for command overhead.
+const reservedCommandBufferBytes = 16 * 10 * 10 * 10
+
 // Insert represents the insert command.
 //
 // The insert command inserts a set of documents into the database.
@@ -27,24 +31,76 @@ type Insert struct {
 	Docs []*bson.Document
 	Opts []option.InsertOptioner
 
-	result result.Insert
-	err    error
+	result          result.Insert
+	err             error
+	continueOnError bool
 }
 
-// Encode will encode this command into a wire message for the given server description.
-func (i *Insert) Encode(desc description.SelectedServer) (wiremessage.WireMessage, error) {
+func (i *Insert) split(maxCount, targetBatchSize int) ([][]*bson.Document, error) {
+	batches := [][]*bson.Document{}
+
+	if targetBatchSize > reservedCommandBufferBytes {
+		targetBatchSize -= reservedCommandBufferBytes
+	}
+
+	if maxCount <= 0 {
+		maxCount = 1
+	}
+
+	startAt := 0
+splitInserts:
+	for {
+		size := 0
+		batch := []*bson.Document{}
+	assembleBatch:
+		for idx := startAt; idx < len(i.Docs); idx++ {
+			itsize, err := i.Docs[idx].Validate()
+			if err != nil {
+				return nil, err
+			}
+
+			if size+int(itsize) > targetBatchSize {
+				break assembleBatch
+			}
+
+			size += int(itsize)
+			batch = append(batch, i.Docs[idx])
+			startAt++
+			if len(batch) == maxCount {
+				break assembleBatch
+			}
+		}
+		batches = append(batches, batch)
+		if startAt == len(i.Docs) {
+			break splitInserts
+		}
+	}
+
+	return batches, nil
+}
+
+func (i *Insert) encodeBatch(docs []*bson.Document, desc description.SelectedServer) (wiremessage.WireMessage, error) {
+
 	command := bson.NewDocument(bson.EC.String("insert", i.NS.Collection))
-	vals := make([]*bson.Value, 0, len(i.Docs))
-	for _, doc := range i.Docs {
+
+	vals := make([]*bson.Value, 0, len(docs))
+	for _, doc := range docs {
 		vals = append(vals, bson.VC.Document(doc))
 	}
 	command.Append(bson.EC.ArrayFromElements("documents", vals...))
 
-	for _, option := range i.Opts {
-		if option == nil {
+	for _, opt := range i.Opts {
+		if opt == nil {
 			continue
 		}
-		err := option.Option(command)
+
+		if ordered, ok := opt.(option.OptOrdered); ok {
+			if !ordered {
+				i.continueOnError = true
+			}
+		}
+
+		err := opt.Option(command)
 		if err != nil {
 			return nil, err
 		}
@@ -53,6 +109,26 @@ func (i *Insert) Encode(desc description.SelectedServer) (wiremessage.WireMessag
 	return (&Command{DB: i.NS.DB, Command: command, isWrite: true}).Encode(desc)
 }
 
+// Encode will encode this command into a wire message for the given server description.
+func (i *Insert) Encode(desc description.SelectedServer) ([]wiremessage.WireMessage, error) {
+	out := []wiremessage.WireMessage{}
+	batches, err := i.split(int(desc.MaxBatchCount), int(desc.MaxDocumentSize))
+	if err != nil {
+		return nil, err
+	}
+
+	for _, docs := range batches {
+		cmd, err := i.encodeBatch(docs, desc)
+		if err != nil {
+			return nil, err
+		}
+
+		out = append(out, cmd)
+	}
+
+	return out, nil
+}
+
 // Decode will decode the wire message using the provided server description. Errors during decoding
 // are deferred until either the Result or Err methods are called.
 func (i *Insert) Decode(desc description.SelectedServer, wm wiremessage.WireMessage) *Insert {
@@ -79,18 +155,40 @@ 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) {
-	wm, err := i.Encode(desc)
-	if err != nil {
-		return result.Insert{}, err
-	}
+	res := result.Insert{}
 
-	err = rw.WriteWireMessage(ctx, wm)
+	wms, err := i.Encode(desc)
 	if err != nil {
-		return result.Insert{}, err
+		return res, err
 	}
-	wm, err = rw.ReadWireMessage(ctx)
-	if err != nil {
-		return result.Insert{}, err
+
+	for _, wm := range wms {
+		err = rw.WriteWireMessage(ctx, wm)
+		if err != nil {
+			return res, err
+		}
+		wm, err = rw.ReadWireMessage(ctx)
+		if err != nil {
+			return res, err
+		}
+
+		r, err := i.Decode(desc, wm).Result()
+		if err != nil {
+			return res, err
+		}
+
+		res.WriteErrors = append(res.WriteErrors, r.WriteErrors...)
+
+		if r.WriteConcernError != nil {
+			res.WriteConcernError = r.WriteConcernError
+		}
+
+		if !i.continueOnError && len(res.WriteErrors) > 0 {
+			return res, nil
+		}
+
+		res.N += r.N
 	}
-	return i.Decode(desc, wm).Result()
+
+	return res, nil
 }

core/command/insert_test.go

@@ -0,0 +1,69 @@
+package command
+
+import (
+	"testing"
+
+	"github.com/mongodb/mongo-go-driver/bson"
+	"github.com/mongodb/mongo-go-driver/core/description"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInsertCommandSplitting(t *testing.T) {
+	const (
+		megabyte = 10 * 10 * 10 * 10 * 10 * 10
+		kilobyte = 10 * 10 * 10
+	)
+
+	ss := description.SelectedServer{}
+	t.Run("split_smoke_test", func(t *testing.T) {
+		i := &Insert{}
+		for n := 0; n < 100; n++ {
+			i.Docs = append(i.Docs, bson.NewDocument(bson.EC.Int32("a", int32(n))))
+		}
+
+		batches, err := i.split(10, kilobyte) // 1kb
+		assert.NoError(t, err)
+		assert.Len(t, batches, 10)
+		for _, b := range batches {
+			assert.Len(t, b, 10)
+			wm, err := i.encodeBatch(b, ss)
+			assert.NoError(t, err)
+			assert.True(t, wm.Len() < 16*megabyte)
+		}
+	})
+	t.Run("split_with_small_target_Size", func(t *testing.T) {
+		i := &Insert{}
+		for n := 0; n < 100; n++ {
+			i.Docs = append(i.Docs, bson.NewDocument(bson.EC.Int32("a", int32(n))))
+		}
+
+		batches, err := i.split(100, 32) // 32 bytes?
+		assert.NoError(t, err)
+		assert.Len(t, batches, 50)
+		for _, b := range batches {
+			assert.Len(t, b, 2)
+			wm, err := i.encodeBatch(b, ss)
+			assert.NoError(t, err)
+			assert.True(t, wm.Len() < 16*megabyte)
+		}
+	})
+	t.Run("invalid_max_counts", func(t *testing.T) {
+		i := &Insert{}
+		for n := 0; n < 100; n++ {
+			i.Docs = append(i.Docs, bson.NewDocument(bson.EC.Int32("a", int32(n))))
+		}
+
+		for _, ct := range []int{-1, 0, -1000} {
+			batches, err := i.split(ct, 100*megabyte)
+			assert.NoError(t, err)
+			assert.Len(t, batches, 100)
+			for _, b := range batches {
+				assert.Len(t, b, 1)
+				wm, err := i.encodeBatch(b, ss)
+				assert.NoError(t, err)
+				assert.True(t, wm.Len() < 16*megabyte)
+			}
+		}
+
+	})
+}

core/description/server.go

@@ -39,7 +39,7 @@ type Server struct {
 	LastError         error
 	LastUpdateTime    time.Time
 	LastWriteTime     time.Time
-	MaxBatchCount     uint16
+	MaxBatchCount     uint32
 	MaxDocumentSize   uint32
 	MaxMessageSize    uint32
 	Members           []address.Address

core/result/result.go

@@ -90,7 +90,7 @@ type IsMaster struct {
 	LastWriteTimestamp  time.Time         `bson:"lastWriteDate,omitempty"`
 	MaxBSONObjectSize   uint32            `bson:"maxBsonObjectSize,omitempty"`
 	MaxMessageSizeBytes uint32            `bson:"maxMessageSizeBytes,omitempty"`
-	MaxWriteBatchSize   uint16            `bson:"maxWriteBatchSize,omitempty"`
+	MaxWriteBatchSize   uint32            `bson:"maxWriteBatchSize,omitempty"`
 	Me                  string            `bson:"me,omitempty"`
 	MaxWireVersion      int32             `bson:"maxWireVersion,omitempty"`
 	MinWireVersion      int32             `bson:"minWireVersion,omitempty"`