Datadog Flush (#379)

.github/workflows/gha-go-test.yaml

@@ -14,6 +14,7 @@ jobs:
         env:
           SC_VERSION: 2023.1.7
         run: |          
+          make generate
           SC_URL="https://github.com/dominikh/go-tools/releases/download/$SC_VERSION/staticcheck_linux_amd64.tar.gz"
           wget -q ${SC_URL} -O - | tar -xzf - --strip-components 1 -C /usr/local/bin staticcheck/staticcheck
           make static

.gitignore

@@ -2,6 +2,9 @@
 # Local .terraform directories
 **/.terraform/*
 
+# Ignore mocks
+**/*.mock.go
+
 # Terraform lock
 *.terraform.lock.hcl
 

lib/kafkalib/writer.go

@@ -141,7 +141,10 @@ func (b *BatchWriter) Write(ctx context.Context, rawMsgs []lib.RawMessage) error
 			}
 		}
 
-		b.statsD.Count("kafka.publish", int64(len(batch)), tags)
+		if b.statsD != nil {
+			b.statsD.Count("kafka.publish", int64(len(batch)), tags)
+		}
+
 		if kafkaErr != nil {
 			return fmt.Errorf("failed to write message: %w, approxSize: %d", kafkaErr, size.GetApproxSize(batch))
 		}

lib/mtr/datadog.go

@@ -18,6 +18,7 @@ type Client interface {
 	Incr(name string, tags map[string]string)
 	Gauge(name string, value float64, tags map[string]string)
 	Count(name string, value int64, tags map[string]string)
+	Flush()
 }
 
 type statsClient struct {
@@ -34,6 +35,10 @@ func toDatadogTags(tags map[string]string) []string {
 	return retTags
 }
 
+func (s *statsClient) Flush() {
+	_ = s.client.Flush()
+}
+
 func (s *statsClient) Count(name string, value int64, tags map[string]string) {
 	_ = s.client.Count(name, value, toDatadogTags(tags), s.rate)
 }

main.go

@@ -6,8 +6,6 @@ import (
 	"fmt"
 	"log/slog"
 
-	"github.com/artie-labs/transfer/lib/telemetry/metrics"
-
 	"github.com/artie-labs/reader/config"
 	"github.com/artie-labs/reader/lib/kafkalib"
 	"github.com/artie-labs/reader/lib/logger"
@@ -23,7 +21,7 @@ import (
 
 func setUpMetrics(cfg *config.Metrics) (mtr.Client, error) {
 	if cfg == nil {
-		return &metrics.NullMetricsProvider{}, nil
+		return nil, nil
 	}
 
 	slog.Info("Creating metrics client")
@@ -86,15 +84,21 @@ func main() {
 	}
 
 	_logger, cleanUpHandlers := logger.NewLogger(cfg)
-	defer cleanUpHandlers()
 	slog.SetDefault(_logger)
-	ctx := context.Background()
 
 	statsD, err := setUpMetrics(cfg.Metrics)
 	if err != nil {
 		logger.Fatal("Failed to set up metrics", slog.Any("err", err))
 	}
 
+	defer func() {
+		cleanUpHandlers()
+		if statsD != nil {
+			statsD.Flush()
+		}
+	}()
+
+	ctx := context.Background()
 	destinationWriter, err := buildDestinationWriter(ctx, cfg, statsD)
 	if err != nil {
 		logger.Fatal(fmt.Sprintf("Failed to init %q destination writer", cfg.Destination), slog.Any("err", err))

writers/transfer/writer.go

@@ -107,7 +107,9 @@ func (w *Writer) Write(_ context.Context, messages []lib.RawMessage) error {
 		"table":    events[0].Table,
 	}
 	defer func() {
-		w.statsD.Count("process.message", int64(len(events)), tags)
+		if w.statsD != nil {
+			w.statsD.Count("process.message", int64(len(events)), tags)
+		}
 	}()
 
 	for _, evt := range events {
@@ -167,7 +169,9 @@ func (w *Writer) flush(reason string) error {
 		"reason":   reason,
 	}
 	defer func() {
-		w.statsD.Timing("flush", time.Since(start), tags)
+		if w.statsD != nil {
+			w.statsD.Timing("flush", time.Since(start), tags)
+		}
 	}()
 
 	if !w.tc.SoftDelete {