Refactor schema caching to use a composite key for SchemaId and format (#2436)

Author: Havret

src/Confluent.SchemaRegistry/CachedSchemaRegistryClient.cs

@@ -58,6 +58,8 @@ namespace Confluent.SchemaRegistry
     /// </summary>
     public class CachedSchemaRegistryClient : ISchemaRegistryClient
     {
+        private record struct SchemaId(int Id, string Format);
+        
         private readonly List<SchemaReference> EmptyReferencesList = new List<SchemaReference>();
 
         private IEnumerable<KeyValuePair<string, string>> config;
@@ -67,7 +69,7 @@ public class CachedSchemaRegistryClient : ISchemaRegistryClient
         private IRestService restService;
         private int identityMapCapacity;
         private int latestCacheTtlSecs;
-        private readonly ConcurrentDictionary<int, Schema> schemaById = new ConcurrentDictionary<int, Schema>();
+        private readonly ConcurrentDictionary<SchemaId, Schema> schemaById = new ConcurrentDictionary<SchemaId, Schema>();
 
         private readonly ConcurrentDictionary<string /*subject*/, ConcurrentDictionary<Schema, int>> idBySchemaBySubject =
             new ConcurrentDictionary<string, ConcurrentDictionary<Schema, int>>();
@@ -632,7 +634,9 @@ public async Task<int> GetSchemaIdAsync(string subject, Schema schema, bool norm
                     var registeredSchema = await restService.LookupSchemaAsync(subject, schema, true, normalize)
                         .ConfigureAwait(continueOnCapturedContext: false);
                     idBySchema[schema] = registeredSchema.Id;
-                    schemaById[registeredSchema.Id] = registeredSchema.Schema;
+                    
+                    var format = GetSchemaFormat(schema.SchemaString);
+                    schemaById.TryAdd(new SchemaId(registeredSchema.Id, format), registeredSchema.Schema);
                     schemaId = registeredSchema.Id;
                 }
 
@@ -695,26 +699,11 @@ public Task<int> RegisterSchemaAsync(string subject, string avroSchema, bool nor
         /// <summary>
         ///     Check if the given schema string matches a given format name.
         /// </summary>
-        private bool checkSchemaMatchesFormat(string format, string schemaString)
+        private static string GetSchemaFormat(string schemaString)
         {
-            // if a format isn't specified, then assume text is desired.
-            if (format == null)
-            {
-                // If schemaString is not Base64, infer the schemaString format is text.
-                return !Utils.IsBase64String(schemaString);
-            }
-            else
-            {
-                if (format != "serialized")
-                {
-                    throw new ArgumentException($"Invalid schema format was specified: {format}.");
-                }
-                
-                return Utils.IsBase64String(schemaString);
-            }
+            return Utils.IsBase64String(schemaString) ? "serialized" : null;
         }
 
-
         /// <inheritdoc/>
         public async Task<RegisteredSchema> LookupSchemaAsync(string subject, Schema schema, bool ignoreDeletedSchemas,
             bool normalize = false)
@@ -753,21 +742,20 @@ public async Task<RegisteredSchema> LookupSchemaAsync(string subject, Schema sch
         /// <inheritdoc/>
         public async Task<Schema> GetSchemaAsync(int id, string format = null)
         {
-            if (schemaById.TryGetValue(id, out var schema) && checkSchemaMatchesFormat(format, schema.SchemaString))
+            var schemaId = new SchemaId(id, format);
+            if (schemaById.TryGetValue(schemaId, out var schema))
             {
                 return schema;
             }
 
             await cacheMutex.WaitAsync().ConfigureAwait(continueOnCapturedContext: false);
             try
             {
-                if (!this.schemaById.TryGetValue(id, out schema) ||
-                    !checkSchemaMatchesFormat(format, schema.SchemaString))
+                if (!this.schemaById.TryGetValue(schemaId, out schema))
                 {
                     CleanCacheIfFull();
-                    schema = (await restService.GetSchemaAsync(id, format)
-                        .ConfigureAwait(continueOnCapturedContext: false));
-                    schemaById[id] = schema;
+                    schema = await restService.GetSchemaAsync(id, format).ConfigureAwait(continueOnCapturedContext: false);
+                    schemaById.TryAdd(schemaId, schema);
                 }
 
                 return schema;
@@ -782,20 +770,21 @@ public async Task<Schema> GetSchemaAsync(int id, string format = null)
         /// <inheritdoc/>
         public async Task<Schema> GetSchemaBySubjectAndIdAsync(string subject, int id, string format = null)
         {
-            if (this.schemaById.TryGetValue(id, out var schema) && checkSchemaMatchesFormat(format, schema.SchemaString))
+            var schemaId = new SchemaId(id, format);
+            if (this.schemaById.TryGetValue(schemaId, out var schema))
             {
                 return schema;
             }
 
             await cacheMutex.WaitAsync().ConfigureAwait(continueOnCapturedContext: false);
             try
             {
-                if (!this.schemaById.TryGetValue(id, out schema) || !checkSchemaMatchesFormat(format, schema.SchemaString))
+                if (!this.schemaById.TryGetValue(schemaId, out schema))
                 {
                     CleanCacheIfFull();
-                    schema = (await restService.GetSchemaBySubjectAndIdAsync(subject, id, format)
-                        .ConfigureAwait(continueOnCapturedContext: false));
-                    schemaById[id] = schema;
+                    schema = await restService.GetSchemaBySubjectAndIdAsync(subject, id, format)
+                        .ConfigureAwait(continueOnCapturedContext: false);
+                    schemaById.TryAdd(schemaId, schema);
                 }
 
                 return schema;
@@ -832,7 +821,8 @@ public async Task<RegisteredSchema> GetRegisteredSchemaAsync(string subject, int
                     schema = await restService.GetSchemaAsync(subject, version)
                         .ConfigureAwait(continueOnCapturedContext: false);
                     schemaByVersion[version] = schema;
-                    schemaById[schema.Id] = schema.Schema;
+                    var format = GetSchemaFormat(schema.SchemaString);
+                    schemaById.TryAdd(new SchemaId(schema.Id, format), schema.Schema);
                 }
 
                 return schema;