CDAP-20488: Cache CDAP schema objects to improve performance (#234)

* CDAP-20488: Cache schema objects to improve performance

* Make schema mapping cache non static

Author: sumitjnn

delta-plugins-common/src/main/java/io/cdap/delta/plugin/common/Records.java

@@ -35,6 +35,8 @@
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.math.BigDecimal;
 import java.time.Instant;
@@ -55,6 +57,8 @@
  * Utilities for converting Records and Schemas.
  */
 public class Records {
+  private static final Logger LOG = LoggerFactory.getLogger(Records.class);
+
   private static final String PRECISION_NAME = "connect.decimal.precision";
   private static final String SCALE_NAME = "scale";
 
@@ -149,10 +153,20 @@ public static StructuredRecord keepSelectedColumns(StructuredRecord record, Set<
    * @param struct
    * @return
    */
-  public static StructuredRecord convert(Struct struct) {
-    Schema schema = convert(struct.schema());
-    schema = schema.isNullable() ? schema.getNonNullable() : schema;
+  public static StructuredRecord convert(Struct struct, SchemaMappingCache schemaMappingCache) {
+    org.apache.kafka.connect.data.Schema schema = struct.schema();
+    Schema mappedSchema = schemaMappingCache.get(schema);
+    if (mappedSchema == null) {
+      LOG.info("Creating CDAP schema from source DB schema");
+      mappedSchema = convert(struct.schema());
+      mappedSchema = mappedSchema.isNullable() ? mappedSchema.getNonNullable() : mappedSchema;
+      schemaMappingCache.put(schema, mappedSchema);
+    }
+    return getStructuredRecord(struct, mappedSchema, schemaMappingCache);
+  }
 
+  private static StructuredRecord getStructuredRecord(Struct struct, Schema schema,
+                                                      SchemaMappingCache schemaMappingCache) {
     StructuredRecord.Builder builder = StructuredRecord.builder(schema);
     if (schema.getFields() == null) {
       return builder.build();
@@ -162,7 +176,7 @@ public static StructuredRecord convert(Struct struct) {
       String fieldName = field.getName();
       Field debeziumField = struct.schema().field(fieldName);
       String debeziumSchemaName = debeziumField.schema().name();
-      Object val = convert(debeziumField.schema(), struct.get(fieldName));
+      Object val = convert(debeziumField.schema(), struct.get(fieldName), schemaMappingCache);
       Schema fieldSchema = field.getSchema();
       fieldSchema = fieldSchema.isNullable() ? fieldSchema.getNonNullable() : fieldSchema;
       Schema.LogicalType logicalType = fieldSchema.getLogicalType();
@@ -257,7 +271,8 @@ private static ZonedDateTime getZonedDateTime(long epochSecond, long fraction, T
   }
 
 
-  private static Object convert(org.apache.kafka.connect.data.Schema schema, Object val) {
+  private static Object convert(org.apache.kafka.connect.data.Schema schema, Object val,
+                                SchemaMappingCache schemaMappingCache) {
     if (val == null) {
       return null;
     }
@@ -275,15 +290,15 @@ private static Object convert(org.apache.kafka.connect.data.Schema schema, Objec
         return ((Short) val).intValue();
       case ARRAY:
         return ((List<?>) val).stream()
-          .map(o -> convert(schema.valueSchema(), o))
+          .map(o -> convert(schema.valueSchema(), o, schemaMappingCache))
           .collect(Collectors.toList());
       case MAP:
         return ((Map<?, ?>) val).entrySet().stream()
           .collect(Collectors.toMap(
-            mapKey -> convert(schema.keySchema(), mapKey),
-            mapVal -> convert(schema.valueSchema(), mapVal)));
+            mapKey -> convert(schema.keySchema(), mapKey, schemaMappingCache),
+            mapVal -> convert(schema.valueSchema(), mapVal, schemaMappingCache)));
       case STRUCT:
-        return convert((Struct) val);
+        return convert((Struct) val, schemaMappingCache);
     }
     // should never happen, all values are listed above
     throw new IllegalStateException(String.format("Kafka type '%s' is not supported.", schema.type()));

delta-plugins-common/src/main/java/io/cdap/delta/plugin/common/SchemaMappingCache.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap.delta.plugin.common;
+
+import com.google.common.collect.MapMaker;
+import io.cdap.cdap.api.data.schema.Schema;
+
+import java.util.Map;
+import javax.annotation.concurrent.ThreadSafe;
+
+/**
+ * Cache which stores mapping of Debezium schema to CDAP schema
+ * Uses reference equality (==) for keys for performance reasons as Schema objects are immutable
+ * Weak references are used  for keys to ensure entries are cleaned up by GC when no longer used by the program
+ */
+@ThreadSafe
+public class SchemaMappingCache {
+
+  private final Map<org.apache.kafka.connect.data.Schema, Schema> cache = new MapMaker()
+    .weakKeys().makeMap();
+
+
+  public void reset() {
+    cache.clear();
+  }
+
+  public void put(org.apache.kafka.connect.data.Schema key, Schema mappedSchema) {
+    cache.put(key, mappedSchema);
+  }
+
+  public Schema get(org.apache.kafka.connect.data.Schema key) {
+    return cache.get(key);
+  }
+}

delta-plugins-common/src/test/java/io/cdap/delta/plugin/common/RecordsTest.java

@@ -44,7 +44,7 @@ public void testConvertTinyInt() {
     Struct struct = new Struct(dataSchema);
     Short val = 1;
     struct.put(fieldName, val);
-    StructuredRecord convert = Records.convert(struct);
+    StructuredRecord convert = Records.convert(struct, new SchemaMappingCache());
     io.cdap.cdap.api.data.schema.Schema.Field priority = convert.getSchema().getField(fieldName);
     Assert.assertNotNull(priority);
     Assert.assertEquals(priority.getSchema(), io.cdap.cdap.api.data.schema.Schema.of(
@@ -61,7 +61,7 @@ public void testConvertTimeStamp() {
     Struct struct = new Struct(dataSchema);
     String val = "2011-12-03T10:15:30.030431+01:00";
     struct.put(fieldName, val);
-    StructuredRecord converted = Records.convert(struct);
+    StructuredRecord converted = Records.convert(struct, new SchemaMappingCache());
     Assert.assertEquals(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS),
                         converted.getSchema().getField(fieldName).getSchema());
     Assert.assertEquals(converted.getTimestamp(fieldName),

mysql-delta-plugins/src/main/java/io/cdap/delta/mysql/MySqlRecordConsumer.java

@@ -26,6 +26,7 @@
 import io.cdap.delta.api.Offset;
 import io.cdap.delta.api.SourceTable;
 import io.cdap.delta.plugin.common.Records;
+import io.cdap.delta.plugin.common.SchemaMappingCache;
 import io.debezium.connector.mysql.MySqlValueConverters;
 import io.debezium.embedded.StopConnectorException;
 import io.debezium.relational.Table;
@@ -49,6 +50,7 @@
  */
 public class MySqlRecordConsumer implements Consumer<SourceRecord> {
   private static final Logger LOG = LoggerFactory.getLogger(MySqlRecordConsumer.class);
+  private static final String TRX_ID_SEP = ":";
 
   private final DeltaSourceContext context;
   private final EventEmitter emitter;
@@ -57,6 +59,7 @@ public class MySqlRecordConsumer implements Consumer<SourceRecord> {
   private final Tables tables;
   private final Map<String, SourceTable> sourceTableMap;
   private final boolean replicateExistingData;
+  private final SchemaMappingCache schemaMappingCache;
 
   public MySqlRecordConsumer(DeltaSourceContext context, EventEmitter emitter,
                              DdlParser ddlParser, MySqlValueConverters mySqlValueConverters,
@@ -68,6 +71,7 @@ public MySqlRecordConsumer(DeltaSourceContext context, EventEmitter emitter,
     this.tables = tables;
     this.sourceTableMap = sourceTableMap;
     this.replicateExistingData = replicateExistingData;
+    this.schemaMappingCache = new SchemaMappingCache();
   }
 
   @Override
@@ -123,7 +127,7 @@ public void accept(SourceRecord sourceRecord) {
     Map<String, String> deltaOffset = generateCdapOffsets(sourceRecord);
     Offset recordOffset = new Offset(deltaOffset);
 
-    StructuredRecord val = Records.convert((Struct) sourceRecord.value());
+    StructuredRecord val = Records.convert((Struct) sourceRecord.value(), schemaMappingCache);
     String ddl = val.get("ddl");
     StructuredRecord source = val.get("source");
     if (source == null) {
@@ -184,9 +188,8 @@ private void handleDML(StructuredRecord source, StructuredRecord val, Offset rec
     String transactionId = source.get("gtid");
     if (transactionId == null) {
       // this is not really a transaction id, but we don't get an event when a transaction started/ended
-      transactionId = String.format("%s:%d",
-                                    source.get(MySqlConstantOffsetBackingStore.FILE),
-                                    source.get(MySqlConstantOffsetBackingStore.POS));
+      transactionId = source.get(MySqlConstantOffsetBackingStore.FILE) + TRX_ID_SEP +
+                                    source.get(MySqlConstantOffsetBackingStore.POS);
     }
 
     StructuredRecord before = val.get("before");
@@ -240,6 +243,7 @@ private void handleDDL(String ddlStatement, Offset recordOffset,
         // CREATE_TABLE and TRUNCATE_TABLE.
         switch (event.type()) {
           case ALTER_TABLE:
+            schemaMappingCache.reset();
             DdlParserListener.TableAlteredEvent alteredEvent = (DdlParserListener.TableAlteredEvent) event;
             TableId tableId = alteredEvent.tableId();
             Table table = tables.forTable(tableId);

sqlserver-delta-plugins/src/main/java/io/cdap/delta/sqlserver/SqlServerRecordConsumer.java

@@ -28,6 +28,7 @@
 import io.cdap.delta.api.Offset;
 import io.cdap.delta.api.SourceTable;
 import io.cdap.delta.plugin.common.Records;
+import io.cdap.delta.plugin.common.SchemaMappingCache;
 import io.debezium.embedded.StopConnectorException;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.source.SourceRecord;
@@ -57,6 +58,7 @@ public class SqlServerRecordConsumer implements Consumer<SourceRecord> {
   private final Map<String, SourceTable> sourceTableMap;
   private final boolean replicateExistingData;
   private final Offset latestOffset;
+  private final SchemaMappingCache schemaMappingCache;
 
 
   SqlServerRecordConsumer(DeltaSourceContext context, EventEmitter emitter, String databaseName,
@@ -69,6 +71,7 @@ public class SqlServerRecordConsumer implements Consumer<SourceRecord> {
     this.sourceTableMap = sourceTableMap;
     this.latestOffset = latestOffset;
     this.replicateExistingData = replicateExistingData;
+    this.schemaMappingCache = new SchemaMappingCache();
   }
 
   @Override
@@ -93,7 +96,7 @@ public void accept(SourceRecord sourceRecord) {
       return;
     }
 
-    StructuredRecord val = Records.convert((Struct) sourceRecord.value());
+    StructuredRecord val = Records.convert((Struct) sourceRecord.value(), schemaMappingCache);
     DMLOperation.Type op;
     String opStr = val.get("op");
     if ("c".equals(opStr) || "r".equals(opStr)) {
@@ -158,7 +161,7 @@ public void accept(SourceRecord sourceRecord) {
         .setSnapshot(ddlRecordOffset.isSnapshot())
         .setOffset(ddlRecordOffset.getAsOffset());
 
-      StructuredRecord key = Records.convert((Struct) sourceRecord.key());
+      StructuredRecord key = Records.convert((Struct) sourceRecord.key(), schemaMappingCache);
       List<Schema.Field> fields = key.getSchema().getFields();
       List<String> primaryFields = new ArrayList<>();
       if (fields != null && !fields.isEmpty()) {