Fix: Implement concurrency control for custom property updates in TypeRepository (#25961)

* Fix: Implement concurrency control for custom property updates in TypeRepository

* Move tests to IT

---------

Co-authored-by: Ram Narayan Balaji <81347100+yan-3005@users.noreply.github.com>

Author: aniketkatkar97

openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/TypeResourceIT.java

@@ -7,9 +7,15 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
@@ -542,6 +548,65 @@ void test_addMultipleHyperlinkCustomProperties(TestNamespace ns) throws Exceptio
     assertTrue(hasProp2, "Type should have second hyperlink custom property");
   }
 
+  @Test
+  void test_concurrentCustomPropertyAdditions(TestNamespace ns) throws Exception {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Type pipelineType = getTypeByName(client, "pipeline");
+
+    int threadCount = 5;
+    List<CustomProperty> properties = new ArrayList<>();
+    for (int i = 0; i < threadCount; i++) {
+      CustomProperty prop = new CustomProperty();
+      prop.setName(ns.prefix("concurrentProp" + i));
+      prop.setDescription("Concurrent property " + i);
+      prop.setPropertyType(STRING_TYPE.getEntityReference());
+      properties.add(prop);
+    }
+
+    ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+    CountDownLatch startLatch = new CountDownLatch(1);
+    CountDownLatch doneLatch = new CountDownLatch(threadCount);
+    List<Exception> errors = new CopyOnWriteArrayList<>();
+
+    for (CustomProperty prop : properties) {
+      executor.submit(
+          () -> {
+            try {
+              startLatch.await();
+              addCustomProperty(client, pipelineType.getId(), prop);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+              errors.add(e);
+            } catch (Exception e) {
+              errors.add(e);
+            } finally {
+              doneLatch.countDown();
+            }
+          });
+    }
+
+    startLatch.countDown();
+    assertTrue(doneLatch.await(30, TimeUnit.SECONDS), "All threads should complete within 30s");
+    executor.shutdown();
+
+    assertTrue(errors.isEmpty(), "Concurrent requests should not throw: " + errors);
+
+    Type updatedType = getTypeByName(client, "pipeline", "customProperties");
+    List<String> persistedNames =
+        updatedType.getCustomProperties() != null
+            ? updatedType.getCustomProperties().stream().map(CustomProperty::getName).toList()
+            : List.of();
+
+    for (CustomProperty prop : properties) {
+      assertTrue(
+          persistedNames.contains(prop.getName()),
+          "Property '"
+              + prop.getName()
+              + "' was lost due to a concurrent update. Persisted: "
+              + persistedNames);
+    }
+  }
+
   private static Type createType(OpenMetadataClient client, CreateType createRequest)
       throws Exception {
     return client
@@ -565,6 +630,16 @@ private static Type getTypeByName(OpenMetadataClient client, String name) throws
     return OBJECT_MAPPER.readValue(response, Type.class);
   }
 
+  private static Type getTypeByName(OpenMetadataClient client, String name, String fields)
+      throws Exception {
+    String response =
+        client
+            .getHttpClient()
+            .executeForString(
+                HttpMethod.GET, "/v1/metadata/types/name/" + name + "?fields=" + fields, null);
+    return OBJECT_MAPPER.readValue(response, Type.class);
+  }
+
   private static TypeList listTypes(OpenMetadataClient client) throws Exception {
     String response =
         client.getHttpClient().executeForString(HttpMethod.GET, "/v1/metadata/types", null);

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TypeRepository.java

@@ -33,6 +33,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.tuple.Triple;
@@ -64,6 +65,8 @@
 public class TypeRepository extends EntityRepository<Type> {
   private static final String UPDATE_FIELDS = "customProperties";
   private static final String PATCH_FIELDS = "customProperties";
+  private static final ConcurrentHashMap<UUID, Object> TYPE_PROPERTY_LOCKS =
+      new ConcurrentHashMap<>();
 
   public TypeRepository() {
     super(
@@ -162,31 +165,34 @@ public void postUpdate(Type original, Type updated) {
 
   public PutResponse<Type> addCustomProperty(
       UriInfo uriInfo, String updatedBy, UUID id, CustomProperty property) {
-    Type type = find(id, Include.NON_DELETED);
-    property.setPropertyType(
-        Entity.getEntityReferenceById(
-            Entity.TYPE, property.getPropertyType().getId(), NON_DELETED));
-    validateProperty(property);
-    if (type.getCategory().equals(Category.Field)) {
-      throw new IllegalArgumentException(
-          "Only entity types can be extended and field types can't be extended");
-    }
-    setFieldsInternal(type, putFields);
+    Object lock = TYPE_PROPERTY_LOCKS.computeIfAbsent(id, k -> new Object());
+    synchronized (lock) {
+      Type type = find(id, Include.NON_DELETED);
+      property.setPropertyType(
+          Entity.getEntityReferenceById(
+              Entity.TYPE, property.getPropertyType().getId(), NON_DELETED));
+      validateProperty(property);
+      if (type.getCategory().equals(Category.Field)) {
+        throw new IllegalArgumentException(
+            "Only entity types can be extended and field types can't be extended");
+      }
+      setFieldsInternal(type, putFields);
 
-    find(property.getPropertyType().getId(), NON_DELETED); // Validate customProperty type exists
+      find(property.getPropertyType().getId(), NON_DELETED); // Validate customProperty type exists
 
-    // If property already exists, then update it. Else add the new property.
-    List<CustomProperty> updatedProperties = new ArrayList<>(List.of(property));
-    for (CustomProperty existing : type.getCustomProperties()) {
-      if (!existing.getName().equals(property.getName())) {
-        updatedProperties.add(existing);
+      // If property already exists, then update it. Else add the new property.
+      List<CustomProperty> updatedProperties = new ArrayList<>(List.of(property));
+      for (CustomProperty existing : type.getCustomProperties()) {
+        if (!existing.getName().equals(property.getName())) {
+          updatedProperties.add(existing);
+        }
       }
-    }
 
-    type.setCustomProperties(updatedProperties);
-    type.setUpdatedBy(updatedBy);
-    type.setUpdatedAt(System.currentTimeMillis());
-    return createOrUpdate(uriInfo, type, updatedBy);
+      type.setCustomProperties(updatedProperties);
+      type.setUpdatedBy(updatedBy);
+      type.setUpdatedAt(System.currentTimeMillis());
+      return createOrUpdate(uriInfo, type, updatedBy);
+    }
   }
 
   private List<CustomProperty> getCustomProperties(Type type) {