[DO NOT MERGE] Microbatch

destination.go

@@ -28,10 +28,14 @@ type Destination struct {
 
 	producer destination.Producer
 	config   destination.Config
+
+	batch MicroBatch
 }
 
 func NewDestination() sdk.Destination {
-	return sdk.DestinationWithMiddleware(&Destination{}, sdk.DefaultDestinationMiddleware()...)
+	return sdk.DestinationWithMiddleware(&Destination{
+		batch: make(MicroBatch, 0, microBatchSize),
+	}, sdk.DefaultDestinationMiddleware()...)
 }
 
 func (d *Destination) Parameters() map[string]sdk.Parameter {
@@ -77,7 +81,13 @@ func (d *Destination) Open(ctx context.Context) error {
 }
 
 func (d *Destination) Write(ctx context.Context, records []sdk.Record) (int, error) {
-	return d.producer.Produce(ctx, records)
+	for _, rec := range records {
+		d.batch.FromRecord(rec)
+	}
+	defer func() {
+		d.batch = d.batch[:0] // reset batch
+	}()
+	return d.producer.Produce(ctx, d.batch)
 }
 
 // Teardown shuts down the Kafka client.

microbatch.go

@@ -0,0 +1,139 @@
+// Copyright © 2023 Meroxa, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// TODO do not push this to main, this is only an experiment to see if micro-batches improve performance
+
+package kafka
+
+import (
+	"bytes"
+	"strconv"
+	"strings"
+
+	sdk "github.com/conduitio/conduit-connector-sdk"
+)
+
+const (
+	separator      = '⌘'
+	microBatchSize = 50
+)
+
+type MicroBatch []sdk.Record
+
+// -- MERGE --
+
+func (b MicroBatch) ToRecord() sdk.Record {
+	lastRec := b[len(b)-1]
+	return sdk.Record{
+		Position:  b.mergePositions(),
+		Operation: lastRec.Operation, // operation is the same for all right now
+		Metadata:  b.mergeMetadata(),
+		Key:       b.mergeKeys(),
+		Payload: sdk.Change{
+			After: b.mergePayloadAfter(),
+		},
+	}
+}
+
+func (b MicroBatch) mergeMetadata() sdk.Metadata {
+	m := make(sdk.Metadata, len(b[0].Metadata)*(len(b)+1))
+	for i, rec := range b {
+		for k, v := range rec.Metadata {
+			m[strconv.Itoa(i)+"."+k] = v
+		}
+	}
+	return m
+}
+
+func (b MicroBatch) mergePositions() sdk.Position {
+	buf := bytes.NewBuffer(make([]byte, 0, len(b[0].Position)*(len(b)+1)))
+	for _, rec := range b {
+		buf.Write(rec.Position)
+		buf.WriteRune(separator)
+	}
+	return buf.Bytes()
+}
+
+func (b MicroBatch) mergeKeys() sdk.Data {
+	buf := bytes.NewBuffer(make([]byte, 0, len(b[0].Key.Bytes())*(len(b)+1)))
+	for _, rec := range b {
+		buf.Write(rec.Key.Bytes())
+		buf.WriteRune(separator)
+	}
+	return sdk.RawData(buf.Bytes())
+}
+
+func (b MicroBatch) mergePayloadAfter() sdk.Data {
+	buf := bytes.NewBuffer(make([]byte, 0, len(b[0].Payload.After.Bytes())*(len(b)+1)))
+	for _, rec := range b {
+		buf.Write(rec.Payload.After.Bytes())
+		buf.WriteRune(separator)
+	}
+	return sdk.RawData(buf.Bytes())
+}
+
+// -- SPLIT --
+
+func (b *MicroBatch) FromRecord(rec sdk.Record) {
+	positions := b.splitBytes(rec.Position)
+	keys := b.splitBytes(rec.Key.Bytes())
+	payloadAfters := b.splitBytes(rec.Payload.After.Bytes())
+	metadata := b.splitMetadata(rec.Metadata, len(positions))
+
+	for i := 0; i < len(positions); i++ {
+		*b = append(*b, sdk.Record{
+			Position:  positions[i],
+			Operation: rec.Operation,
+			Metadata:  metadata[i],
+			Key:       sdk.RawData(keys[i]),
+			Payload: sdk.Change{
+				After: sdk.RawData(payloadAfters[i]),
+			},
+		})
+	}
+}
+
+func (MicroBatch) splitBytes(bb []byte) [][]byte {
+	split := bytes.Split(bb, []byte(string([]rune{separator})))
+	return split[:len(split)-1] // data ends with a separator, remove last element
+}
+
+func (MicroBatch) splitMetadata(metadata sdk.Metadata, count int) []sdk.Metadata {
+	m := make([]sdk.Metadata, count)
+	genericMetadata := make(sdk.Metadata)
+	for k, v := range metadata {
+		prefix, key, found := strings.Cut(k, ".")
+		if !found {
+			genericMetadata[k] = v
+			continue
+		}
+		i, err := strconv.Atoi(prefix)
+		if err != nil {
+			genericMetadata[k] = v
+			continue
+		}
+		md := m[i]
+		if md == nil {
+			md = sdk.Metadata{}
+			m[i] = md
+		}
+		md[key] = v
+	}
+	for _, md := range m {
+		for k, v := range genericMetadata {
+			md[k] = v
+		}
+	}
+	return m
+}

microbatch_test.go

@@ -0,0 +1,55 @@
+// Copyright © 2023 Meroxa, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package kafka
+
+import (
+	"fmt"
+	"testing"
+
+	sdk "github.com/conduitio/conduit-connector-sdk"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/google/uuid"
+	"github.com/matryer/is"
+)
+
+func TestMicroBatch(t *testing.T) {
+	is := is.New(t)
+
+	have := make(MicroBatch, 10)
+	for i := range have {
+		have[i] = sdk.Record{
+			Position:  sdk.Position(fmt.Sprintf("pos-%d", i)),
+			Operation: sdk.OperationCreate,
+			Metadata: sdk.Metadata{
+				"foo":                    uuid.NewString(),
+				fmt.Sprintf("foo-%d", i): fmt.Sprintf("special-metadata-for-%d", i),
+			},
+			Key: sdk.RawData(fmt.Sprintf("key-%d", i)),
+			Payload: sdk.Change{
+				After: sdk.RawData(fmt.Sprintf("after-%d", i)),
+			},
+		}
+	}
+
+	// merge into one record
+	r := have.ToRecord()
+
+	// split batch back into multiple records
+	got := make(MicroBatch, 0, len(have))
+	got.FromRecord(r)
+
+	is.Equal(cmp.Diff(got, have, cmpopts.IgnoreUnexported(sdk.Record{})), "")
+}

source.go

@@ -33,10 +33,14 @@ type Source struct {
 
 	consumer source.Consumer
 	config   source.Config
+
+	batch MicroBatch
 }
 
 func NewSource() sdk.Source {
-	return sdk.SourceWithMiddleware(&Source{}, sdk.DefaultSourceMiddleware()...)
+	return sdk.SourceWithMiddleware(&Source{
+		batch: make(MicroBatch, microBatchSize),
+	}, sdk.DefaultSourceMiddleware()...)
 }
 
 func (s *Source) Parameters() map[string]sdk.Parameter {
@@ -91,29 +95,39 @@ func (s *Source) Open(ctx context.Context, sdkPos sdk.Position) error {
 }
 
 func (s *Source) Read(ctx context.Context) (sdk.Record, error) {
-	rec, err := s.consumer.Consume(ctx)
-	if err != nil {
-		return sdk.Record{}, fmt.Errorf("failed getting a record: %w", err)
+	for i := 0; i < cap(s.batch); i++ {
+		rec, err := s.consumer.Consume(ctx)
+		if err != nil {
+			return sdk.Record{}, fmt.Errorf("failed getting a record: %w", err)
+		}
+
+		metadata := sdk.Metadata{MetadataKafkaTopic: rec.Topic}
+		metadata.SetCreatedAt(rec.Timestamp)
+
+		s.batch[i] = sdk.Util.Source.NewRecordCreate(
+			source.Position{
+				GroupID:   s.config.GroupID,
+				Topic:     rec.Topic,
+				Partition: rec.Partition,
+				Offset:    rec.Offset,
+			}.ToSDKPosition(),
+			metadata,
+			sdk.RawData(rec.Key),
+			sdk.RawData(rec.Value),
+		)
 	}
 
-	metadata := sdk.Metadata{MetadataKafkaTopic: rec.Topic}
-	metadata.SetCreatedAt(rec.Timestamp)
-
-	return sdk.Util.Source.NewRecordCreate(
-		source.Position{
-			GroupID:   s.config.GroupID,
-			Topic:     rec.Topic,
-			Partition: rec.Partition,
-			Offset:    rec.Offset,
-		}.ToSDKPosition(),
-		metadata,
-		sdk.RawData(rec.Key),
-		sdk.RawData(rec.Value),
-	), nil
+	return s.batch.ToRecord(), nil
 }
 
 func (s *Source) Ack(ctx context.Context, _ sdk.Position) error {
-	return s.consumer.Ack(ctx)
+	for i := 0; i < microBatchSize; i++ {
+		err := s.consumer.Ack(ctx)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
 func (s *Source) Teardown(ctx context.Context) error {