Proto schema registry refresh optimizations

backend/pkg/schema/client.go

@@ -211,19 +211,19 @@ func (c *Client) GetSchemaBySubject(ctx context.Context, subject, version string
 }
 
 // GetSchemasBySubject gets all versioned schemas for a given subject.
-func (c *Client) GetSchemasBySubject(ctx context.Context, subject string, showSoftDeleted bool) ([]SchemaVersionedResponse, error) {
+func (c *Client) GetSchemasBySubject(ctx context.Context, subject string, showSoftDeleted bool) ([]*SchemaVersionedResponse, error) {
 	versionRes, err := c.GetSubjectVersions(ctx, subject, showSoftDeleted)
 	if err != nil {
 		return nil, err
 	}
 
-	results := []SchemaVersionedResponse{}
+	results := []*SchemaVersionedResponse{}
 	for _, sv := range versionRes.Versions {
 		sr, err := c.GetSchemaBySubject(ctx, subject, strconv.FormatInt(int64(sv), 10), showSoftDeleted)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get schema by subject %s and version %d", subject, sv)
 		}
-		results = append(results, *sr)
+		results = append(results, sr)
 	}
 
 	return results, nil
@@ -595,8 +595,8 @@ func (c *Client) GetSchemaTypes(ctx context.Context) ([]string, error) {
 }
 
 // GetSchemas retrieves all stored schemas from a schema registry.
-func (c *Client) GetSchemas(ctx context.Context, showSoftDeleted bool) ([]SchemaVersionedResponse, error) {
-	var schemas []SchemaVersionedResponse
+func (c *Client) GetSchemas(ctx context.Context, showSoftDeleted bool) ([]*SchemaVersionedResponse, []error) {
+	var schemas []*SchemaVersionedResponse
 	req := c.client.R().
 		SetContext(ctx).
 		SetResult(&schemas)
@@ -607,7 +607,7 @@ func (c *Client) GetSchemas(ctx context.Context, showSoftDeleted bool) ([]Schema
 
 	res, err := req.Get("/schemas")
 	if err != nil {
-		return nil, fmt.Errorf("get schemas failed: %w", err)
+		return nil, []error{fmt.Errorf("get schemas failed: %w", err)}
 	}
 
 	if res.StatusCode() == http.StatusNotFound {
@@ -619,9 +619,9 @@ func (c *Client) GetSchemas(ctx context.Context, showSoftDeleted bool) ([]Schema
 	if res.IsError() {
 		restErr, ok := res.Error().(*RestError)
 		if !ok {
-			return nil, fmt.Errorf("get schemas failed: Status code %d", res.StatusCode())
+			return nil, []error{fmt.Errorf("get schemas failed: Status code %d", res.StatusCode())}
 		}
-		return nil, restErr
+		return nil, []error{restErr}
 	}
 
 	return schemas, nil
@@ -687,38 +687,44 @@ func (c *Client) CreateSchema(ctx context.Context, subjectName string, schema Sc
 
 // GetSchemasIndividually returns all schemas by describing all schemas one by one. This may be used against
 // schema registry that don't support the /schemas endpoint that returns a list of all registered schemas.
-func (c *Client) GetSchemasIndividually(ctx context.Context, showSoftDeleted bool) ([]SchemaVersionedResponse, error) {
+func (c *Client) GetSchemasIndividually(ctx context.Context, showSoftDeleted bool) ([]*SchemaVersionedResponse, []error) {
 	subjectsRes, err := c.GetSubjects(ctx, showSoftDeleted)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get subjects to fetch schemas for: %w", err)
+		return nil, []error{fmt.Errorf("failed to get subjects to fetch schemas for: %w", err)}
 	}
 
 	g, ctx := errgroup.WithContext(ctx)
 	g.SetLimit(10) // limit max concurrency to 10 requests at a time
 
-	schemas := make([]SchemaVersionedResponse, 0, len(subjectsRes.Subjects))
+	schemas := make([]*SchemaVersionedResponse, 0, len(subjectsRes.Subjects))
+	errors := make([]error, 0, len(subjectsRes.Subjects))
 	mutex := sync.Mutex{}
 
 	for _, subject := range subjectsRes.Subjects {
 		subject := subject
+
 		g.Go(func() error {
 			srRes, err := c.GetSchemasBySubject(ctx, subject, showSoftDeleted)
 			if err != nil {
-				return err
+				mutex.Lock()
+				errors = append(errors, err)
+				mutex.Unlock()
+				return nil
 			}
 
 			mutex.Lock()
 			schemas = append(schemas, srRes...)
 			mutex.Unlock()
-			return err
+			return nil
 		})
 	}
 
+	// is this even possible?
 	if err := g.Wait(); err != nil {
-		return nil, err
+		return nil, []error{err}
 	}
 
-	return schemas, nil
+	return schemas, errors
 }
 
 // GetSchemaReferencesResponse is the response to fetching schema references.

backend/pkg/schema/service.go

@@ -12,8 +12,10 @@ package schema
 import (
 	"context"
 	"fmt"
+	"maps"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/hamba/avro/v2"
@@ -41,6 +43,11 @@ type Service struct {
 	// by subjects is needed to lookup references in avro schemas.
 	schemaBySubjectVersion *cache.Cache[string, *SchemaVersionedResponse]
 	avroSchemaByID         *cache.Cache[uint32, avro.Schema]
+
+	// for protobuf schema refreshing and compiling
+	srRefreshMutex   sync.RWMutex
+	protoSchemasByID map[int]*SchemaVersionedResponse
+	protoFDByID      map[int]*desc.FileDescriptor
 }
 
 // NewService to access schema registry. Returns an error if connection can't be established.
@@ -67,35 +74,73 @@ func (s *Service) CheckConnectivity(ctx context.Context) error {
 
 // GetProtoDescriptors returns all file descriptors in a map where the key is the schema id.
 // The value is a set of file descriptors because each schema may references / imported proto schemas.
+//
+//nolint:gocognit,cyclop // complicated refresh logic
 func (s *Service) GetProtoDescriptors(ctx context.Context) (map[int]*desc.FileDescriptor, error) {
 	// Singleflight makes sure to not run the function body if there are concurrent requests. We use this to avoid
 	// duplicate requests against the schema registry
 	key := "get-proto-descriptors"
-	v, err, _ := s.requestGroup.Do(key, func() (any, error) {
-		schemasRes, err := s.registryClient.GetSchemas(ctx, false)
-		if err != nil {
-			// If schema registry returns an error we want to retry it next time, so let's forget the key
-			s.requestGroup.Forget(key)
-			return nil, fmt.Errorf("failed to get schema from registry: %w", err)
+	_, err, _ := s.requestGroup.Do(key, func() (any, error) {
+		schemasRes, errs := s.registryClient.GetSchemas(ctx, false)
+		if len(errs) > 0 {
+			for _, err := range errs {
+				s.logger.Error("failed to get schema from registry", zap.Error(err))
+			}
+
+			if len(schemasRes) == 0 {
+				return nil, nil
+			}
+		}
+
+		s.srRefreshMutex.Lock()
+		defer s.srRefreshMutex.Unlock()
+
+		if s.protoSchemasByID == nil {
+			s.protoSchemasByID = make(map[int]*SchemaVersionedResponse, len(schemasRes))
+		}
+
+		// collect existing schema IDs
+		existingSchemaIDs := make(map[int]struct{}, len(s.protoSchemasByID))
+		for id := range s.protoSchemasByID {
+			existingSchemaIDs[id] = struct{}{}
 		}
 
-		// 1. Index all returned schemas by their respective subject name and version as stored in the schema registry
-		schemasBySubjectAndVersion := make(map[string]map[int]SchemaVersionedResponse)
+		schemasToCompile := make([]*SchemaVersionedResponse, 0, len(schemasRes))
+
+		newSchemaIDs := make(map[int]struct{}, len(schemasRes))
+
+		// Index all returned schemas by their respective subject name and version as stored in the schema registry
+		// Collect the new or updated schemas to compile
+		schemasBySubjectAndVersion := make(map[string]map[int]*SchemaVersionedResponse)
 		for _, schema := range schemasRes {
+			schema := schema
+
 			if schema.Type != TypeProtobuf {
 				continue
 			}
 			_, exists := schemasBySubjectAndVersion[schema.Subject]
 			if !exists {
-				schemasBySubjectAndVersion[schema.Subject] = make(map[int]SchemaVersionedResponse)
+				schemasBySubjectAndVersion[schema.Subject] = make(map[int]*SchemaVersionedResponse)
 			}
 			schemasBySubjectAndVersion[schema.Subject][schema.Version] = schema
+
+			if existing, ok := s.protoSchemasByID[schema.SchemaID]; !ok || !strings.EqualFold(existing.Schema, schema.Schema) {
+				schemasToCompile = append(schemasToCompile, schema)
+			}
+
+			newSchemaIDs[schema.SchemaID] = struct{}{}
+
+			s.protoSchemasByID[schema.SchemaID] = schema
 		}
 
-		// 2. Compile each subject with each of it's references into one or more filedescriptors so that they can be
+		compileStart := time.Now()
+
+		// 2. Compile each subject with each of it's references into one or more file descriptors so that they can be
 		// registered in their own proto registry.
-		fdBySchemaID := make(map[int]*desc.FileDescriptor)
-		for _, schema := range schemasRes {
+		newFDBySchemaID := make(map[int]*desc.FileDescriptor, len(schemasToCompile))
+		for _, schema := range schemasToCompile {
+			schema := schema
+
 			if schema.Type != TypeProtobuf {
 				continue
 			}
@@ -108,24 +153,57 @@ func (s *Service) GetProtoDescriptors(ctx context.Context) (map[int]*desc.FileDe
 					zap.Error(err))
 				continue
 			}
-			fdBySchemaID[schema.SchemaID] = fd
+			newFDBySchemaID[schema.SchemaID] = fd
+		}
+
+		compileDuration := time.Since(compileStart)
+
+		// merge
+		if s.protoFDByID == nil {
+			s.protoFDByID = make(map[int]*desc.FileDescriptor, len(newFDBySchemaID))
+		}
+
+		maps.Copy(s.protoFDByID, newFDBySchemaID)
+
+		schemasDeleted := 0
+
+		// remove schemas only if no errors
+		if len(errs) == 0 {
+			schemasIDsToDelete := make([]int, 0, len(s.protoSchemasByID)/2)
+
+			for id := range existingSchemaIDs {
+				if _, ok := newSchemaIDs[id]; !ok {
+					schemasIDsToDelete = append(schemasIDsToDelete, id)
+				}
+			}
+
+			for _, id := range schemasIDsToDelete {
+				delete(s.protoSchemasByID, id)
+				delete(s.protoFDByID, id)
+			}
+
+			schemasDeleted = len(schemasIDsToDelete)
 		}
 
-		return fdBySchemaID, nil
+		s.logger.Info("compiled new schemas",
+			zap.Int("updated_schemas", len(schemasToCompile)),
+			zap.Int("deleted_schemas", schemasDeleted),
+			zap.Duration("compile_duration", compileDuration))
+
+		return nil, nil
 	})
 	if err != nil {
 		return nil, err
 	}
 
-	descriptors, ok := v.(map[int]*desc.FileDescriptor)
-	if !ok {
-		return nil, fmt.Errorf("failed to type assert file descriptors")
-	}
+	s.srRefreshMutex.RLock()
+	descriptors := maps.Clone(s.protoFDByID)
+	s.srRefreshMutex.RUnlock()
 
 	return descriptors, nil
 }
 
-func (s *Service) addReferences(schema SchemaVersionedResponse, schemaRepository map[string]map[int]SchemaVersionedResponse, schemasByPath map[string]string) error {
+func (s *Service) addReferences(schema *SchemaVersionedResponse, schemaRepository map[string]map[int]*SchemaVersionedResponse, schemasByPath map[string]string) error {
 	for _, ref := range schema.References {
 		refSubject, exists := schemaRepository[ref.Subject]
 		if !exists {
@@ -147,7 +225,7 @@ func (s *Service) addReferences(schema SchemaVersionedResponse, schemaRepository
 	return nil
 }
 
-func (s *Service) compileProtoSchemas(schema SchemaVersionedResponse, schemaRepository map[string]map[int]SchemaVersionedResponse) (*desc.FileDescriptor, error) {
+func (s *Service) compileProtoSchemas(schema *SchemaVersionedResponse, schemaRepository map[string]map[int]*SchemaVersionedResponse) (*desc.FileDescriptor, error) {
 	// 1. Let's find the references for each schema and put the references' schemas into our in memory filesystem.
 	schemasByPath := make(map[string]string)
 	schemasByPath[schema.Subject] = schema.Schema

taskfiles/backend.yaml

@@ -37,7 +37,7 @@ tasks:
 
   install-gofumpt:
     vars:
-      GOFUMPT_VERSION: 0.4.0
+      GOFUMPT_VERSION: 0.5.0
     desc: install gofumpt go formater
     deps:
       - install-go
@@ -50,7 +50,7 @@ tasks:
 
   install-gci:
     vars:
-      GCI_VERSION: 0.9.0
+      GCI_VERSION: 0.12.1
     desc: install gci
     deps:
       - install-go
@@ -108,7 +108,9 @@ tasks:
     cmds:
       - '{{ .BUILD_ROOT }}/bin/go/goimports -l -w -local "github.com/redpanda-data/console/backend" .'
       - '{{ .BUILD_ROOT }}/bin/go/gofumpt -l -w .'
-      - '{{ .BUILD_ROOT }}/bin/go/gci write -s standard -s default -s ''Prefix(github.com/redpanda-data/console/backend)'' .'
+      # https://github.com/daixiang0/gci/issues/174
+      # moves "maps" into the wrong place
+      # - '{{ .BUILD_ROOT }}/bin/go/gci write -s default -s standard -s ''prefix(github.com/redpanda-data/console/backend)'' .'
 
   lint:
     desc: Run Go linters for backend code