add optional caching mode to serializer (#1151)

fzmoment · web-flow · commit f1230c0e9dd4 · 2024-03-15T14:48:44.000-07:00
diff --git a/schemaregistry/serde/protobuf/config.go b/schemaregistry/serde/protobuf/config.go
@@ -21,12 +21,21 @@ import "github.com/confluentinc/confluent-kafka-go/v2/schemaregistry/serde"
 // SerializerConfig is used to pass multiple configuration options to the serializers.
 type SerializerConfig struct {
 	serde.SerializerConfig
+	// CacheSchemas will cache serialization results based on the name of the protobuf file
+	// corresponding to the message being serialized. This will drastically improve serialization
+	// performance if you are only ever using a _single_ version of a specific protobuf schema
+	// during any given run of your application. This should be the case for most applications,
+	// but might not apply if you're not creating proto messages based on generated files (e.g.
+	// you are proxying or reading raw protobuf messages from a data source), or if for some reason
+	// you are including multiple versions of the same schema/protobuf in your application.
+	CacheSchemas bool
 }
 
 // NewSerializerConfig returns a new configuration instance with sane defaults.
 func NewSerializerConfig() *SerializerConfig {
 	c := &SerializerConfig{
 		SerializerConfig: *serde.NewSerializerConfig(),
+		CacheSchemas:     false,
 	}
 
 	return c
diff --git a/schemaregistry/serde/protobuf/protobuf.go b/schemaregistry/serde/protobuf/protobuf.go
@@ -66,6 +66,9 @@ import (
 // Serializer represents a Protobuf serializer
 type Serializer struct {
 	serde.BaseSerializer
+	Conf                  *SerializerConfig
+	descToSchemaCache     cache.Cache
+	descToSchemaCacheLock sync.RWMutex
 }
 
 // Deserializer represents a Protobuf deserializer
@@ -132,8 +135,16 @@ func init() {
 
 // NewSerializer creates a Protobuf serializer for Protobuf-generated objects
 func NewSerializer(client schemaregistry.Client, serdeType serde.Type, conf *SerializerConfig) (*Serializer, error) {
-	s := &Serializer{}
-	err := s.ConfigureSerializer(client, serdeType, &conf.SerializerConfig)
+	cache, err := cache.NewLRUCache(1000)
+	if err != nil {
+		return nil, err
+	}
+	s := &Serializer{
+		descToSchemaCache: cache,
+	}
+	err = s.ConfigureSerializer(client, serdeType, &conf.SerializerConfig)
+	s.Conf = conf
+
 	if err != nil {
 		return nil, err
 	}
@@ -166,22 +177,11 @@ func (s *Serializer) Serialize(topic string, msg interface{}) ([]byte, error) {
 	default:
 		return nil, fmt.Errorf("serialization target must be a protobuf message. Got '%v'", t)
 	}
-	autoRegister := s.Conf.AutoRegisterSchemas
-	normalize := s.Conf.NormalizeSchemas
-	fileDesc, deps, err := s.toProtobufSchema(protoMsg)
+	info, err := s.getSchemaInfo(protoMsg)
 	if err != nil {
 		return nil, err
 	}
-	metadata, err := s.resolveDependencies(fileDesc, deps, "", autoRegister, normalize)
-	if err != nil {
-		return nil, err
-	}
-	info := schemaregistry.SchemaInfo{
-		Schema:     metadata.Schema,
-		SchemaType: metadata.SchemaType,
-		References: metadata.References,
-	}
-	id, err := s.GetID(topic, protoMsg, &info)
+	id, err := s.GetID(topic, protoMsg, info)
 	if err != nil {
 		return nil, err
 	}
@@ -197,18 +197,50 @@ func (s *Serializer) Serialize(topic string, msg interface{}) ([]byte, error) {
 	return payload, nil
 }
 
-func (s *Serializer) toProtobufSchema(msg proto.Message) (*desc.FileDescriptor, map[string]string, error) {
-	messageDesc, err := desc.LoadMessageDescriptorForMessage(protoV1.MessageV1(msg))
+func (s *Serializer) getSchemaInfo(protoMsg proto.Message) (*schemaregistry.SchemaInfo, error) {
+	messageDesc, err := desc.LoadMessageDescriptorForMessage(protoV1.MessageV1(protoMsg))
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 	fileDesc := messageDesc.GetFile()
+	if s.Conf.CacheSchemas {
+		s.descToSchemaCacheLock.RLock()
+		value, ok := s.descToSchemaCache.Get(fileDesc.GetName())
+		s.descToSchemaCacheLock.RUnlock()
+		if ok {
+			return value.(*schemaregistry.SchemaInfo), nil
+		}
+	}
+	deps, err := s.toProtobufSchema(fileDesc)
+	if err != nil {
+		return nil, err
+	}
+	autoRegister := s.Conf.AutoRegisterSchemas
+	normalize := s.Conf.NormalizeSchemas
+	metadata, err := s.resolveDependencies(fileDesc, deps, "", autoRegister, normalize)
+	if err != nil {
+		return nil, err
+	}
+	info := &schemaregistry.SchemaInfo{
+		Schema:     metadata.Schema,
+		SchemaType: metadata.SchemaType,
+		References: metadata.References,
+	}
+	if s.Conf.CacheSchemas {
+		s.descToSchemaCacheLock.Lock()
+		s.descToSchemaCache.Put(fileDesc.GetName(), info)
+		s.descToSchemaCacheLock.Unlock()
+	}
+	return info, nil
+}
+
+func (s *Serializer) toProtobufSchema(fileDesc *desc.FileDescriptor) (map[string]string, error) {
 	deps := make(map[string]string)
-	err = s.toDependencies(fileDesc, deps)
+	err := s.toDependencies(fileDesc, deps)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
-	return fileDesc, deps, nil
+	return deps, nil
 }
 
 func (s *Serializer) toDependencies(fileDesc *desc.FileDescriptor, deps map[string]string) error {
diff --git a/schemaregistry/serde/protobuf/protobuf_test.go b/schemaregistry/serde/protobuf/protobuf_test.go
@@ -199,3 +199,130 @@ func TestProtobufSerdeEmptyMessage(t *testing.T) {
 	_, err = deser.Deserialize("topic1", []byte{})
 	serde.MaybeFail("deserialization", err)
 }
+
+func BenchmarkProtobufSerWithReference(b *testing.B) {
+	serde.MaybeFail = serde.InitFailFuncBenchmark(b)
+	var err error
+	conf := schemaregistry.NewConfig("mock://")
+
+	client, err := schemaregistry.NewClient(conf)
+	serde.MaybeFail("Schema Registry configuration", err)
+
+	ser, err := NewSerializer(client, serde.ValueSerde, NewSerializerConfig())
+	serde.MaybeFail("Serializer configuration", err)
+
+	msg := test.TestMessage{
+		TestString:   "hi",
+		TestBool:     true,
+		TestBytes:    []byte{1, 2},
+		TestDouble:   1.23,
+		TestFloat:    3.45,
+		TestFixed32:  67,
+		TestFixed64:  89,
+		TestInt32:    100,
+		TestInt64:    200,
+		TestSfixed32: 300,
+		TestSfixed64: 400,
+		TestSint32:   500,
+		TestSint64:   600,
+		TestUint32:   700,
+		TestUint64:   800,
+	}
+	obj := test.DependencyMessage{
+		IsActive:     true,
+		TestMesssage: &msg,
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ser.Serialize("topic1", &obj)
+	}
+}
+
+func BenchmarkProtobufSerWithReferenceCached(b *testing.B) {
+	serde.MaybeFail = serde.InitFailFuncBenchmark(b)
+	var err error
+	conf := schemaregistry.NewConfig("mock://")
+
+	client, err := schemaregistry.NewClient(conf)
+	serde.MaybeFail("Schema Registry configuration", err)
+
+	serConf := NewSerializerConfig()
+	serConf.CacheSchemas = true
+	ser, err := NewSerializer(client, serde.ValueSerde, serConf)
+	serde.MaybeFail("Serializer configuration", err)
+
+	msg := test.TestMessage{
+		TestString:   "hi",
+		TestBool:     true,
+		TestBytes:    []byte{1, 2},
+		TestDouble:   1.23,
+		TestFloat:    3.45,
+		TestFixed32:  67,
+		TestFixed64:  89,
+		TestInt32:    100,
+		TestInt64:    200,
+		TestSfixed32: 300,
+		TestSfixed64: 400,
+		TestSint32:   500,
+		TestSint64:   600,
+		TestUint32:   700,
+		TestUint64:   800,
+	}
+	obj := test.DependencyMessage{
+		IsActive:     true,
+		TestMesssage: &msg,
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ser.Serialize("topic1", &obj)
+	}
+}
+
+func BenchmarkProtobufDeserWithReference(b *testing.B) {
+	serde.MaybeFail = serde.InitFailFuncBenchmark(b)
+	var err error
+	conf := schemaregistry.NewConfig("mock://")
+
+	client, err := schemaregistry.NewClient(conf)
+	serde.MaybeFail("Schema Registry configuration", err)
+
+	ser, err := NewSerializer(client, serde.ValueSerde, NewSerializerConfig())
+	serde.MaybeFail("Serializer configuration", err)
+
+	msg := test.TestMessage{
+		TestString:   "hi",
+		TestBool:     true,
+		TestBytes:    []byte{1, 2},
+		TestDouble:   1.23,
+		TestFloat:    3.45,
+		TestFixed32:  67,
+		TestFixed64:  89,
+		TestInt32:    100,
+		TestInt64:    200,
+		TestSfixed32: 300,
+		TestSfixed64: 400,
+		TestSint32:   500,
+		TestSint64:   600,
+		TestUint32:   700,
+		TestUint64:   800,
+	}
+	obj := test.DependencyMessage{
+		IsActive:     true,
+		TestMesssage: &msg,
+	}
+	bytes, err := ser.Serialize("topic1", &obj)
+	serde.MaybeFail("serialization", err)
+
+	deser, err := NewDeserializer(client, serde.ValueSerde, NewDeserializerConfig())
+	serde.MaybeFail("Deserializer configuration", err)
+	deser.Client = ser.Client
+
+	deser.ProtoRegistry.RegisterMessage(obj.ProtoReflect().Type())
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		deser.Deserialize("topic1", bytes)
+	}
+}
diff --git a/schemaregistry/serde/testhelpers.go b/schemaregistry/serde/testhelpers.go
@@ -46,6 +46,22 @@ func InitFailFunc(t *testing.T) FailFunc {
 	}
 }
 
+func InitFailFuncBenchmark(b *testing.B) FailFunc {
+	tester := b
+	return func(msg string, errors ...error) {
+		for _, err := range errors {
+			if err != nil {
+				pc := make([]uintptr, 1)
+				runtime.Callers(2, pc)
+				caller := runtime.FuncForPC(pc[0])
+				_, line := caller.FileLine(caller.Entry())
+
+				tester.Fatalf("%s:%d failed: %s %s", caller.Name(), line, msg, err)
+			}
+		}
+	}
+}
+
 // Expect compares the actual and expected values
 func Expect(actual, expected interface{}) error {
 	if !reflect.DeepEqual(actual, expected) {
