Repository: xataio/pgstream
Pull request #85 (2024-10-25 10:05:01)

Wrap external json lib 
This PR wraps the external json lib implementation under an internal package so that it can be easily replaced in all places where it's used. It updates processor side modules to benefit from the performance improvements of the library change.

It also updates the `GetTableByName` method to reduce unnecessary memory allocations.

Follow up on https://github.com/xataio/pgstream/pull/84.


Patch:
diff --git a/internal/json/json.go b/internal/json/json.go
new file mode 100644
index 0000000..37c515a
--- /dev/null
+++ b/internal/json/json.go
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package json
+
+import (
+	json "github.com/bytedance/sonic"
+)
+
+func Unmarshal(b []byte, v any) error {
+	return json.Unmarshal(b, v)
+}
+
+func Marshal(v any) ([]byte, error) {
+	return json.Marshal(v)
+}
diff --git a/internal/searchstore/search_client.go b/internal/searchstore/search_client.go
index 19b7d69..2cc0f67 100644
--- a/internal/searchstore/search_client.go
+++ b/internal/searchstore/search_client.go
@@ -5,9 +5,10 @@ package searchstore
 import (
 	"bytes"
 	"context"
-	"encoding/json"
 	"fmt"
 	"net/http"
+
+	"github.com/xataio/pgstream/internal/json"
 )
 
 type Client interface {
diff --git a/pkg/schemalog/log_entry.go b/pkg/schemalog/log_entry.go
index ae5cb24..b5b7989 100644
--- a/pkg/schemalog/log_entry.go
+++ b/pkg/schemalog/log_entry.go
@@ -96,7 +96,7 @@ func (m *LogEntry) Diff(previous *LogEntry) *SchemaDiff {
 	return m.Schema.Diff(&previous.Schema)
 }
 
-func (m *LogEntry) GetTableByName(tableName string) *Table {
+func (m *LogEntry) GetTableByName(tableName string) (Table, bool) {
 	return m.Schema.getTableByName(tableName)
 }
 
diff --git a/pkg/schemalog/schema.go b/pkg/schemalog/schema.go
index 3f4b700..a913ca2 100644
--- a/pkg/schemalog/schema.go
+++ b/pkg/schemalog/schema.go
@@ -3,8 +3,9 @@
 package schemalog
 
 import (
-	"encoding/json"
 	"slices"
+
+	"github.com/xataio/pgstream/internal/json"
 )
 
 type Schema struct {
@@ -95,13 +96,13 @@ func (s *Schema) Diff(previous *Schema) *SchemaDiff {
 	return &d
 }
 
-func (s *Schema) getTableByName(tableName string) *Table {
+func (s *Schema) getTableByName(tableName string) (Table, bool) {
 	for _, t := range s.Tables {
 		if t.Name == tableName {
-			return &t
+			return t, true
 		}
 	}
-	return nil
+	return Table{}, false
 }
 
 func (s *Schema) getTableByID(pgstreamID string) *Table {
diff --git a/pkg/wal/listener/kafka/wal_kafka_reader.go b/pkg/wal/listener/kafka/wal_kafka_reader.go
index 0954f8b..63927bf 100644
--- a/pkg/wal/listener/kafka/wal_kafka_reader.go
+++ b/pkg/wal/listener/kafka/wal_kafka_reader.go
@@ -4,10 +4,10 @@ package kafka
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/pkg/kafka"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/wal"
diff --git a/pkg/wal/listener/postgres/wal_pg_listener.go b/pkg/wal/listener/postgres/wal_pg_listener.go
index 40b88b0..55d1e37 100644
--- a/pkg/wal/listener/postgres/wal_pg_listener.go
+++ b/pkg/wal/listener/postgres/wal_pg_listener.go
@@ -7,7 +7,7 @@ import (
 	"errors"
 	"fmt"
 
-	json "github.com/bytedance/sonic"
+	"github.com/xataio/pgstream/internal/json"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/wal"
 	"github.com/xataio/pgstream/pkg/wal/replication"
diff --git a/pkg/wal/processor/kafka/wal_kafka_batch_writer.go b/pkg/wal/processor/kafka/wal_kafka_batch_writer.go
index dbfa938..3849d8e 100644
--- a/pkg/wal/processor/kafka/wal_kafka_batch_writer.go
+++ b/pkg/wal/processor/kafka/wal_kafka_batch_writer.go
@@ -9,7 +9,7 @@ import (
 	"runtime/debug"
 	"time"
 
-	json "github.com/bytedance/sonic"
+	"github.com/xataio/pgstream/internal/json"
 	synclib "github.com/xataio/pgstream/internal/sync"
 	"github.com/xataio/pgstream/pkg/kafka"
 	kafkainstrumentation "github.com/xataio/pgstream/pkg/kafka/instrumentation"
diff --git a/pkg/wal/processor/search/search_adapter.go b/pkg/wal/processor/search/search_adapter.go
index 0104851..84af34b 100644
--- a/pkg/wal/processor/search/search_adapter.go
+++ b/pkg/wal/processor/search/search_adapter.go
@@ -3,11 +3,11 @@
 package search
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"strconv"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/wal"
 	"github.com/xataio/pgstream/pkg/wal/processor"
diff --git a/pkg/wal/processor/search/search_store_retrier.go b/pkg/wal/processor/search/search_store_retrier.go
index 2ecccab..549b5b8 100644
--- a/pkg/wal/processor/search/search_store_retrier.go
+++ b/pkg/wal/processor/search/search_store_retrier.go
@@ -4,11 +4,11 @@ package search
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"time"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/pkg/backoff"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/schemalog"
@@ -19,6 +19,7 @@ type StoreRetrier struct {
 	inner           Store
 	logger          loglib.Logger
 	backoffProvider backoff.Provider
+	marshaler       func(any) ([]byte, error)
 }
 
 type StoreRetryConfig struct {
@@ -42,6 +43,7 @@ func NewStoreRetrier(s Store, cfg StoreRetryConfig, opts ...StoreOption) *StoreR
 		inner:           s,
 		logger:          loglib.NewNoopLogger(),
 		backoffProvider: backoff.NewProvider(cfg.backoffConfig()),
+		marshaler:       json.Marshal,
 	}
 
 	for _, opt := range opts {
@@ -171,7 +173,7 @@ func (s *StoreRetrier) getRetriableDocs(failedDocs []DocumentError) ([]Document,
 }
 
 func (s *StoreRetrier) logFailure(docErr DocumentError) {
-	docBytes, err := json.Marshal(docErr.Document.Data)
+	docBytes, err := s.marshaler(docErr.Document.Data)
 	if err != nil {
 		docBytes = []byte(fmt.Sprintf("failed to marshal document data: %s", err))
 	}
diff --git a/pkg/wal/processor/search/store/search_adapter.go b/pkg/wal/processor/search/store/search_adapter.go
index 81cdf95..05ae0ce 100644
--- a/pkg/wal/processor/search/store/search_adapter.go
+++ b/pkg/wal/processor/search/store/search_adapter.go
@@ -3,9 +3,9 @@
 package store
 
 import (
-	"encoding/json"
 	"fmt"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/internal/searchstore"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
diff --git a/pkg/wal/processor/search/store/search_pg_mapper.go b/pkg/wal/processor/search/store/search_pg_mapper.go
index 15d906e..a064dad 100644
--- a/pkg/wal/processor/search/store/search_pg_mapper.go
+++ b/pkg/wal/processor/search/store/search_pg_mapper.go
@@ -3,7 +3,6 @@
 package store
 
 import (
-	"encoding/json"
 	"fmt"
 	"strconv"
 	"strings"
@@ -11,6 +10,7 @@ import (
 
 	"github.com/jackc/pgx/v5/pgtype"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/internal/searchstore"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
@@ -19,6 +19,7 @@ import (
 type PgMapper struct {
 	searchMapper searchstore.Mapper
 	pgTypeMap    *pgtype.Map
+	unmarshaler  func([]byte, any) error
 }
 
 const (
@@ -34,6 +35,7 @@ func NewPostgresMapper(mapper searchstore.Mapper) *PgMapper {
 	return &PgMapper{
 		searchMapper: mapper,
 		pgTypeMap:    pgtype.NewMap(),
+		unmarshaler:  json.Unmarshal,
 	}
 }
 
@@ -82,7 +84,7 @@ func (m *PgMapper) MapColumnValue(column schemalog.Column, value any) (any, erro
 			return nil, fmt.Errorf("unexpected value type for jsonb column")
 		}
 		var array []float64
-		err := json.Unmarshal([]byte(stringContent), &array)
+		err := m.unmarshaler([]byte(stringContent), &array)
 		if err != nil {
 			return nil, fmt.Errorf("vector value is not array: %w", err)
 		}
diff --git a/pkg/wal/processor/search/store/search_store.go b/pkg/wal/processor/search/store/search_store.go
index 6c55dc4..a283da3 100644
--- a/pkg/wal/processor/search/store/search_store.go
+++ b/pkg/wal/processor/search/store/search_store.go
@@ -5,14 +5,13 @@ package store
 import (
 	"bytes"
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 
+	"github.com/xataio/pgstream/internal/json"
 	"github.com/xataio/pgstream/internal/searchstore"
 	elasticsearchstore "github.com/xataio/pgstream/internal/searchstore/elasticsearch"
 	opensearchstore "github.com/xataio/pgstream/internal/searchstore/opensearch"
-
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/schemalog"
 	"github.com/xataio/pgstream/pkg/wal/processor/search"
@@ -423,7 +422,7 @@ func (s *Store) updateMappingAddNewColumns(ctx context.Context, indexName IndexN
 }
 
 func (s *Store) insertNewSchemaLog(ctx context.Context, m *schemalog.LogEntry) error {
-	logBytes, err := json.Marshal(m)
+	logBytes, err := s.marshaler(m)
 	if err != nil {
 		return fmt.Errorf("insert schema log, failed to marshal search doc: %w", err)
 	}
diff --git a/pkg/wal/processor/translator/wal_translator.go b/pkg/wal/processor/translator/wal_translator.go
index 7cfd916..058cf10 100644
--- a/pkg/wal/processor/translator/wal_translator.go
+++ b/pkg/wal/processor/translator/wal_translator.go
@@ -203,16 +203,16 @@ func (t *Translator) translate(ctx context.Context, data *wal.Data) error {
 		return fmt.Errorf("failed to retrieve schema for translate %w", err)
 	}
 
-	table := logEntry.GetTableByName(data.Table)
-	if table == nil {
+	table, found := logEntry.GetTableByName(data.Table)
+	if !found {
 		return processor.ErrTableNotFound
 	}
 
-	if err = t.fillEventMetadata(data, logEntry, table); err != nil {
+	if err = t.fillEventMetadata(data, logEntry, &table); err != nil {
 		return fmt.Errorf("failed to fill event metadata: %w", err)
 	}
 
-	if err = t.translateColumnNames(data, table); err != nil {
+	if err = t.translateColumnNames(data, &table); err != nil {
 		return fmt.Errorf("failed to translate column names: %w", err)
 	}
 
diff --git a/pkg/wal/processor/webhook/notifier/webhook_notifier.go b/pkg/wal/processor/webhook/notifier/webhook_notifier.go
index ac1b4d7..b4d8b0f 100644
--- a/pkg/wal/processor/webhook/notifier/webhook_notifier.go
+++ b/pkg/wal/processor/webhook/notifier/webhook_notifier.go
@@ -5,7 +5,6 @@ package notifier
 import (
 	"bytes"
 	"context"
-	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
@@ -13,6 +12,7 @@ import (
 	"sync"
 
 	httplib "github.com/xataio/pgstream/internal/http"
+	"github.com/xataio/pgstream/internal/json"
 	synclib "github.com/xataio/pgstream/internal/sync"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/wal"


Test patch:
diff --git a/pkg/wal/processor/search/search_store_retrier_test.go b/pkg/wal/processor/search/search_store_retrier_test.go
index d471a1c..c371bf0 100644
--- a/pkg/wal/processor/search/search_store_retrier_test.go
+++ b/pkg/wal/processor/search/search_store_retrier_test.go
@@ -4,6 +4,7 @@ package search
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"testing"
 
@@ -229,6 +230,7 @@ func TestStoreRetrier_SendDocuments(t *testing.T) {
 				inner:           tc.store,
 				logger:          loglib.NewNoopLogger(),
 				backoffProvider: newMockBackoffProvider(),
+				marshaler:       json.Marshal,
 			}
 
 			failedDocs, err := retrier.SendDocuments(context.Background(), testDocs)