Send the events as soon as we have hit the batch size limit (#686)

This is for performance, we want to keep `START_REPLICATION` connection alive for as long as we can. This was observed to reduce the latency overall significantly.
PeerDB-io · Nov 20, 2023 · 2a60ff3 · 2a60ff3
1 parent b4cdb12
commit 2a60ff3
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 86 deletions.
diff --git a/flow/connectors/eventhub/eventhub.go b/flow/connectors/eventhub/eventhub.go
@@ -126,54 +126,70 @@ func (c *EventHubConnector) processBatch(
 	batchPerTopic := NewHubBatches(c.hubManager)
 	toJSONOpts := model.NewToJSONOptions(c.config.UnnestColumns)
 
+	eventHubFlushTimeout :=
+		time.Duration(utils.GetEnvInt("PEERDB_EVENTHUB_FLUSH_TIMEOUT_SECONDS", 10)) *
+			time.Second
+
+	ticker := time.NewTicker(eventHubFlushTimeout)
+	defer ticker.Stop()
+
 	numRecords := 0
-	for record := range batch.GetRecords() {
-		numRecords++
-		json, err := record.GetItems().ToJSONWithOpts(toJSONOpts)
-		if err != nil {
-			log.WithFields(log.Fields{
-				"flowName": flowJobName,
-			}).Infof("failed to convert record to json: %v", err)
-			return 0, err
-		}
+	for {
+		select {
+		case record, ok := <-batch.GetRecords():
+			if !ok {
+				err := batchPerTopic.flushAllBatches(ctx, maxParallelism, flowJobName)
+				if err != nil {
+					return 0, err
+				}
+
+				log.WithFields(log.Fields{
+					"flowName": flowJobName,
+				}).Infof("[total] successfully sent %d records to event hub", numRecords)
+				return uint32(numRecords), nil
+			}
 
-		topicName, err := NewScopedEventhub(record.GetTableName())
-		if err != nil {
-			log.WithFields(log.Fields{
-				"flowName": flowJobName,
-			}).Infof("failed to get topic name: %v", err)
-			return 0, err
-		}
+			numRecords++
+			json, err := record.GetItems().ToJSONWithOpts(toJSONOpts)
+			if err != nil {
+				log.WithFields(log.Fields{
+					"flowName": flowJobName,
+				}).Infof("failed to convert record to json: %v", err)
+				return 0, err
+			}
 
-		err = batchPerTopic.AddEvent(ctx, topicName, json)
-		if err != nil {
-			log.WithFields(log.Fields{
-				"flowName": flowJobName,
-			}).Infof("failed to add event to batch: %v", err)
-			return 0, err
-		}
+			topicName, err := NewScopedEventhub(record.GetTableName())
+			if err != nil {
+				log.WithFields(log.Fields{
+					"flowName": flowJobName,
+				}).Infof("failed to get topic name: %v", err)
+				return 0, err
+			}
 
-		if numRecords%1000 == 0 {
-			log.WithFields(log.Fields{
-				"flowName": flowJobName,
-			}).Infof("processed %d records for sending", numRecords)
-		}
-	}
+			err = batchPerTopic.AddEvent(ctx, topicName, json, false)
+			if err != nil {
+				log.WithFields(log.Fields{
+					"flowName": flowJobName,
+				}).Infof("failed to add event to batch: %v", err)
+				return 0, err
+			}
 
-	log.WithFields(log.Fields{
-		"flowName": flowJobName,
-	}).Infof("processed %d records for sending", numRecords)
+			if numRecords%1000 == 0 {
+				log.WithFields(log.Fields{
+					"flowName": flowJobName,
+				}).Infof("processed %d records for sending", numRecords)
+			}
 
-	flushErr := batchPerTopic.flushAllBatches(ctx, maxParallelism, flowJobName)
-	if flushErr != nil {
-		return 0, flushErr
-	}
-	batchPerTopic.Clear()
+		case <-ticker.C:
+			err := batchPerTopic.flushAllBatches(ctx, maxParallelism, flowJobName)
+			if err != nil {
+				return 0, err
+			}
 
-	log.WithFields(log.Fields{
-		"flowName": flowJobName,
-	}).Infof("[total] successfully flushed %d records to event hub", numRecords)
-	return uint32(numRecords), nil
+			ticker.Stop()
+			ticker = time.NewTicker(eventHubFlushTimeout)
+		}
+	}
 }
 
 func (c *EventHubConnector) SyncRecords(req *model.SyncRecordsRequest) (*model.SyncResponse, error) {

diff --git a/flow/connectors/eventhub/hub_batches.go b/flow/connectors/eventhub/hub_batches.go
@@ -2,8 +2,8 @@ package conneventhub
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"strings"
 	"sync/atomic"
 	"time"
 
@@ -14,72 +14,70 @@ import (
 
 // multimap from ScopedEventhub to *azeventhubs.EventDataBatch
 type HubBatches struct {
-	batches map[ScopedEventhub][]*azeventhubs.EventDataBatch
+	batch   map[ScopedEventhub]*azeventhubs.EventDataBatch
 	manager *EventHubManager
 }
 
 func NewHubBatches(manager *EventHubManager) *HubBatches {
 	return &HubBatches{
-		batches: make(map[ScopedEventhub][]*azeventhubs.EventDataBatch),
+		batch:   make(map[ScopedEventhub]*azeventhubs.EventDataBatch),
 		manager: manager,
 	}
 }
 
-func (h *HubBatches) AddEvent(ctx context.Context, name ScopedEventhub, event string) error {
-	batches, ok := h.batches[name]
-	if !ok {
-		batches = []*azeventhubs.EventDataBatch{}
-	}
-
-	if len(batches) == 0 {
+func (h *HubBatches) AddEvent(
+	ctx context.Context,
+	name ScopedEventhub,
+	event string,
+	// this is true when we are retrying to send the event after the batch size exceeded
+	// this should initially be false, and then true when we are retrying.
+	retryForBatchSizeExceed bool,
+) error {
+	batch, ok := h.batch[name]
+	if !ok || batch == nil {
 		newBatch, err := h.manager.CreateEventDataBatch(ctx, name)
 		if err != nil {
-			return err
+			return fmt.Errorf("failed to create event data batch: %v", err)
 		}
-		batches = append(batches, newBatch)
+		batch = newBatch
+		h.batch[name] = batch
 	}
 
-	if err := tryAddEventToBatch(event, batches[len(batches)-1]); err != nil {
-		if strings.Contains(err.Error(), "too large for the batch") {
-			overflowBatch, err := h.handleBatchOverflow(ctx, name, event)
-			if err != nil {
-				return fmt.Errorf("failed to handle batch overflow: %v", err)
-			}
-			batches = append(batches, overflowBatch)
-		} else {
-			return fmt.Errorf("failed to add event data: %v", err)
-		}
+	err := tryAddEventToBatch(event, batch)
+	if err == nil {
+		// we successfully added the event to the batch, so we're done.
+		return nil
 	}
 
-	h.batches[name] = batches
-	return nil
-}
+	if errors.Is(err, azeventhubs.ErrEventDataTooLarge) {
+		if retryForBatchSizeExceed {
+			// if we are already retrying, then we should just return the error
+			// as we have already tried to send the event to the batch.
+			return fmt.Errorf("[retry-failed] event too large to add to batch: %v", err)
+		}
 
-func (h *HubBatches) handleBatchOverflow(
-	ctx context.Context,
-	name ScopedEventhub,
-	event string,
-) (*azeventhubs.EventDataBatch, error) {
-	newBatch, err := h.manager.CreateEventDataBatch(ctx, name)
-	if err != nil {
-		return nil, err
-	}
-	if err := tryAddEventToBatch(event, newBatch); err != nil {
-		return nil, fmt.Errorf("failed to add event data to new batch: %v", err)
+		// if the event is too large, send the current batch and
+		// delete it from the map, so that a new batch can be created
+		// for the event next time.
+		if err := h.sendBatch(ctx, name, batch); err != nil {
+			return fmt.Errorf("failed to send batch: %v", err)
+		}
+		delete(h.batch, name)
+
+		return h.AddEvent(ctx, name, event, true)
+	} else {
+		return fmt.Errorf("failed to add event to batch: %v", err)
 	}
-	return newBatch, nil
 }
 
 func (h *HubBatches) Len() int {
-	return len(h.batches)
+	return len(h.batch)
 }
 
 // ForEach calls the given function for each ScopedEventhub and batch pair
 func (h *HubBatches) ForEach(fn func(ScopedEventhub, *azeventhubs.EventDataBatch)) {
-	for name, batches := range h.batches {
-		for _, batch := range batches {
-			fn(name, batch)
-		}
+	for name, batch := range h.batch {
+		fn(name, batch)
 	}
 }
 
@@ -137,14 +135,19 @@ func (h *HubBatches) flushAllBatches(
 		})
 	})
 
-	log.Infof("[sendEventBatch] successfully sent %d events in total to event hub",
+	err := g.Wait()
+	log.Infof("[flush] successfully sent %d events in total to event hub",
 		numEventsPushed)
-	return g.Wait()
+
+	// clear the batches after flushing them.
+	h.Clear()
+
+	return err
 }
 
 // Clear removes all batches from the HubBatches
 func (h *HubBatches) Clear() {
-	h.batches = make(map[ScopedEventhub][]*azeventhubs.EventDataBatch)
+	h.batch = make(map[ScopedEventhub]*azeventhubs.EventDataBatch)
 }
 
 func tryAddEventToBatch(event string, batch *azeventhubs.EventDataBatch) error {