From ea779e8b2f217a60c14303ceef88b8a5f4924f93 Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Tue, 23 Nov 2021 00:23:01 +0530
Subject: [PATCH 1/3] Changes to support bulk create of labels

---
 .../service/v1/labels_config_service.proto    | 16 ++++++
 .../service/LabelsConfigServiceImpl.java      | 56 +++++++++++++++++++
 .../service/LabelsConfigServiceImplTest.java  | 51 ++++++++++++++++-
 3 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto b/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
index be57bc20..2154575c 100644
--- a/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
+++ b/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
@@ -11,6 +11,9 @@ service LabelsConfigService {
   // create label
   rpc CreateLabel(CreateLabelRequest) returns (CreateLabelResponse) {}
 
+  // create labels
+  rpc BulkCreateLabels(BulkCreateLabelsRequest) returns (BulkCreateLabelsResponse) {}
+
   // get label by id
   rpc GetLabel(GetLabelRequest) returns (GetLabelResponse) {}
 
@@ -34,6 +37,19 @@ message CreateLabelResponse {
   Label label = 1;
 }
 
+message BulkCreateLabelsRequest {
+  repeated LabelRequest requests = 1;
+
+  message LabelRequest {
+    LabelData data = 1;
+    optional string created_by_application_rule_id = 2;
+  }
+}
+
+message BulkCreateLabelsResponse {
+  repeated Label labels = 1;
+}
+
 message GetLabelRequest {
   string id = 1;
 }
diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index a89e689f..bd9cbc35 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -13,6 +13,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
@@ -25,6 +26,8 @@
 import org.hypertrace.config.service.v1.ConfigServiceGrpc;
 import org.hypertrace.core.grpcutils.client.RequestContextClientCallCredsProviderFactory;
 import org.hypertrace.core.grpcutils.context.RequestContext;
+import org.hypertrace.label.config.service.v1.BulkCreateLabelsRequest;
+import org.hypertrace.label.config.service.v1.BulkCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.CreateLabelRequest;
 import org.hypertrace.label.config.service.v1.CreateLabelResponse;
 import org.hypertrace.label.config.service.v1.DeleteLabelRequest;
@@ -136,6 +139,59 @@ public void createLabel(
     }
   }
 
+  @Override
+  public void bulkCreateLabels(
+      BulkCreateLabelsRequest bulkCreateLabelsRequest,
+      StreamObserver<BulkCreateLabelsResponse> responseObserver) {
+    try {
+      RequestContext requestContext = RequestContext.CURRENT.get();
+      Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
+      if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
+        try {
+          List<Label> labels =
+              bulkCreateLabelsRequest.getRequestsList().stream()
+                  .map(this::buildLabelFromRequest)
+                  .filter(Optional::isPresent)
+                  .map(Optional::get)
+                  .collect(Collectors.toList());
+          if (!labels.isEmpty()) {
+            List<Label> createdLabels =
+                labelStore.upsertObjects(requestContext, labels).stream()
+                    .map(org.hypertrace.config.objectstore.ConfigObject::getData)
+                    .collect(Collectors.toList());
+            responseObserver.onNext(
+                BulkCreateLabelsResponse.newBuilder().addAllLabels(createdLabels).build());
+          } else {
+            responseObserver.onNext(BulkCreateLabelsResponse.newBuilder().build());
+          }
+          responseObserver.onCompleted();
+        } finally {
+          labelsLock.unlock();
+        }
+      } else {
+        responseObserver.onError(new StatusRuntimeException(Status.ABORTED));
+      }
+    } catch (Exception e) {
+      responseObserver.onError(e);
+    }
+  }
+
+  private Optional<Label> buildLabelFromRequest(BulkCreateLabelsRequest.LabelRequest request) {
+    LabelData labelData = request.getData();
+    if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())
+        || isDuplicateKey(labelData.getKey())) {
+      // Creating a label with a name that clashes with one of system labels
+      // name. i
+      return Optional.empty();
+    }
+    Label.Builder labelBuilder =
+        Label.newBuilder().setId(UUID.randomUUID().toString()).setData(labelData);
+    if (request.hasCreatedByApplicationRuleId()) {
+      labelBuilder.setCreatedByApplicationRuleId(request.getCreatedByApplicationRuleId());
+    }
+    return Optional.of(labelBuilder.build());
+  }
+
   @Override
   public void getLabel(GetLabelRequest request, StreamObserver<GetLabelResponse> responseObserver) {
     RequestContext requestContext = RequestContext.CURRENT.get();
diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index f38fe530..a6fa65eb 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -9,6 +9,7 @@
 import io.grpc.Channel;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
@@ -18,6 +19,8 @@
 import java.util.stream.Stream;
 import org.hypertrace.config.service.change.event.api.ConfigChangeEventGenerator;
 import org.hypertrace.config.service.test.MockGenericConfigService;
+import org.hypertrace.label.config.service.v1.BulkCreateLabelsRequest;
+import org.hypertrace.label.config.service.v1.BulkCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.CreateLabelRequest;
 import org.hypertrace.label.config.service.v1.CreateLabelResponse;
 import org.hypertrace.label.config.service.v1.DeleteLabelRequest;
@@ -55,7 +58,12 @@ public final class LabelsConfigServiceImplTest {
   @BeforeEach
   void setUp() {
     mockGenericConfigService =
-        new MockGenericConfigService().mockUpsert().mockGet().mockGetAll().mockDelete();
+        new MockGenericConfigService()
+            .mockUpsert()
+            .mockUpsertAll()
+            .mockGet()
+            .mockGetAll()
+            .mockDelete();
     Map<String, List<Map<String, String>>> systemLabelsConfigMap = new HashMap<>();
     systemLabelsConfigMap.put(
         "system.labels",
@@ -125,6 +133,47 @@ void test_system_createLabel() {
     }
   }
 
+  @Test
+  void test_bulkCreateLabels() {
+    List<BulkCreateLabelsRequest.LabelRequest> requests =
+        createLabelDataList.stream()
+            .map(data -> BulkCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
+            .collect(Collectors.toList());
+    BulkCreateLabelsResponse bulkCreateLabelsResponse =
+        labelConfigStub.bulkCreateLabels(
+            BulkCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
+    List<LabelData> createdLabelsList =
+        bulkCreateLabelsResponse.getLabelsList().stream()
+            .map(label -> label.getData())
+            .collect(Collectors.toList());
+    assertEquals(createLabelDataList, createdLabelsList);
+  }
+
+  @Test
+  void test_bulkCreateLabelsWithSystemLabels() {
+    List<BulkCreateLabelsRequest.LabelRequest> requests = new ArrayList<>();
+    requests.addAll(
+        createLabelDataList.stream()
+            .map(data -> BulkCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
+            .collect(Collectors.toList()));
+    requests.addAll(
+        systemLabels.stream()
+            .map(
+                systemLabel ->
+                    BulkCreateLabelsRequest.LabelRequest.newBuilder()
+                        .setData(systemLabel.getData())
+                        .build())
+            .collect(Collectors.toList()));
+    BulkCreateLabelsResponse bulkCreateLabelsResponse =
+        labelConfigStub.bulkCreateLabels(
+            BulkCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
+    List<LabelData> createdLabelsList =
+        bulkCreateLabelsResponse.getLabelsList().stream()
+            .map(label -> label.getData())
+            .collect(Collectors.toList());
+    assertEquals(createLabelDataList, createdLabelsList);
+  }
+
   @Test
   void test_getLabel() {
     List<Label> createdLabelsList =

From 7720b97e74cadc0497c92c77c00f606dadee3f94 Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Tue, 23 Nov 2021 23:52:53 +0530
Subject: [PATCH 2/3] Address review comments

---
 .../service/v1/labels_config_service.proto    |  8 +-
 .../service/LabelsConfigServiceImpl.java      | 80 ++++++++++---------
 .../service/LabelsConfigServiceImplTest.java  | 62 +++++++++-----
 3 files changed, 87 insertions(+), 63 deletions(-)

diff --git a/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto b/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
index 2154575c..9fad2811 100644
--- a/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
+++ b/labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto
@@ -11,8 +11,8 @@ service LabelsConfigService {
   // create label
   rpc CreateLabel(CreateLabelRequest) returns (CreateLabelResponse) {}
 
-  // create labels
-  rpc BulkCreateLabels(BulkCreateLabelsRequest) returns (BulkCreateLabelsResponse) {}
+  // get or create labels
+  rpc GetOrCreateLabels(GetOrCreateLabelsRequest) returns (GetOrCreateLabelsResponse) {}
 
   // get label by id
   rpc GetLabel(GetLabelRequest) returns (GetLabelResponse) {}
@@ -37,7 +37,7 @@ message CreateLabelResponse {
   Label label = 1;
 }
 
-message BulkCreateLabelsRequest {
+message GetOrCreateLabelsRequest {
   repeated LabelRequest requests = 1;
 
   message LabelRequest {
@@ -46,7 +46,7 @@ message BulkCreateLabelsRequest {
   }
 }
 
-message BulkCreateLabelsResponse {
+message GetOrCreateLabelsResponse {
   repeated Label labels = 1;
 }
 
diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index bd9cbc35..bf58889b 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -1,5 +1,7 @@
 package org.hypertrace.label.config.service;
 
+import static java.util.function.Function.identity;
+
 import com.google.common.util.concurrent.Striped;
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.JsonFormat;
@@ -13,11 +15,9 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
@@ -26,8 +26,6 @@
 import org.hypertrace.config.service.v1.ConfigServiceGrpc;
 import org.hypertrace.core.grpcutils.client.RequestContextClientCallCredsProviderFactory;
 import org.hypertrace.core.grpcutils.context.RequestContext;
-import org.hypertrace.label.config.service.v1.BulkCreateLabelsRequest;
-import org.hypertrace.label.config.service.v1.BulkCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.CreateLabelRequest;
 import org.hypertrace.label.config.service.v1.CreateLabelResponse;
 import org.hypertrace.label.config.service.v1.DeleteLabelRequest;
@@ -36,6 +34,8 @@
 import org.hypertrace.label.config.service.v1.GetLabelResponse;
 import org.hypertrace.label.config.service.v1.GetLabelsRequest;
 import org.hypertrace.label.config.service.v1.GetLabelsResponse;
+import org.hypertrace.label.config.service.v1.GetOrCreateLabelsRequest;
+import org.hypertrace.label.config.service.v1.GetOrCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.Label;
 import org.hypertrace.label.config.service.v1.LabelData;
 import org.hypertrace.label.config.service.v1.LabelsConfigServiceGrpc;
@@ -74,13 +74,11 @@ public LabelsConfigServiceImpl(
     if (systemLabelsObjectList != null) {
       systemLabels = buildSystemLabelList(systemLabelsObjectList);
       systemLabelsIdLabelMap =
-          systemLabels.stream()
-              .collect(Collectors.toUnmodifiableMap(Label::getId, Function.identity()));
+          systemLabels.stream().collect(Collectors.toUnmodifiableMap(Label::getId, identity()));
       systemLabelsKeyLabelMap =
           systemLabels.stream()
               .collect(
-                  Collectors.toUnmodifiableMap(
-                      (label) -> label.getData().getKey(), Function.identity()));
+                  Collectors.toUnmodifiableMap((label) -> label.getData().getKey(), identity()));
     } else {
       systemLabels = Collections.emptyList();
       systemLabelsIdLabelMap = Collections.emptyMap();
@@ -140,30 +138,31 @@ public void createLabel(
   }
 
   @Override
-  public void bulkCreateLabels(
-      BulkCreateLabelsRequest bulkCreateLabelsRequest,
-      StreamObserver<BulkCreateLabelsResponse> responseObserver) {
+  public void getOrCreateLabels(
+      GetOrCreateLabelsRequest request,
+      StreamObserver<GetOrCreateLabelsResponse> responseObserver) {
     try {
       RequestContext requestContext = RequestContext.CURRENT.get();
       Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
       if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
         try {
-          List<Label> labels =
-              bulkCreateLabelsRequest.getRequestsList().stream()
-                  .map(this::buildLabelFromRequest)
-                  .filter(Optional::isPresent)
-                  .map(Optional::get)
-                  .collect(Collectors.toList());
-          if (!labels.isEmpty()) {
+          final Map<String, Label> existingLabelsMap = getLabelsMap();
+          final List<Label> newLabels = new ArrayList<>();
+          final List<Label> allLabels = new ArrayList<>();
+          request.getRequestsList().stream()
+              .forEach(
+                  labelRequest ->
+                      this.buildLabelFromRequest(
+                          labelRequest, existingLabelsMap, newLabels, allLabels));
+          if (!newLabels.isEmpty()) {
             List<Label> createdLabels =
-                labelStore.upsertObjects(requestContext, labels).stream()
+                labelStore.upsertObjects(requestContext, newLabels).stream()
                     .map(org.hypertrace.config.objectstore.ConfigObject::getData)
                     .collect(Collectors.toList());
-            responseObserver.onNext(
-                BulkCreateLabelsResponse.newBuilder().addAllLabels(createdLabels).build());
-          } else {
-            responseObserver.onNext(BulkCreateLabelsResponse.newBuilder().build());
+            allLabels.addAll(createdLabels);
           }
+          responseObserver.onNext(
+              GetOrCreateLabelsResponse.newBuilder().addAllLabels(allLabels).build());
           responseObserver.onCompleted();
         } finally {
           labelsLock.unlock();
@@ -176,20 +175,26 @@ public void bulkCreateLabels(
     }
   }
 
-  private Optional<Label> buildLabelFromRequest(BulkCreateLabelsRequest.LabelRequest request) {
+  private void buildLabelFromRequest(
+      GetOrCreateLabelsRequest.LabelRequest request,
+      Map<String, Label> existingLabelsMap,
+      List<Label> newLabels,
+      List<Label> allLabels) {
     LabelData labelData = request.getData();
-    if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())
-        || isDuplicateKey(labelData.getKey())) {
-      // Creating a label with a name that clashes with one of system labels
-      // name. i
-      return Optional.empty();
+    if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())) {
+      allLabels.add(systemLabelsKeyLabelMap.get(labelData.getKey()));
+      return;
+    }
+    if (existingLabelsMap.containsKey(labelData.getKey())) {
+      allLabels.add(existingLabelsMap.get(labelData.getKey()));
+      return;
     }
     Label.Builder labelBuilder =
         Label.newBuilder().setId(UUID.randomUUID().toString()).setData(labelData);
     if (request.hasCreatedByApplicationRuleId()) {
       labelBuilder.setCreatedByApplicationRuleId(request.getCreatedByApplicationRuleId());
     }
-    return Optional.of(labelBuilder.build());
+    newLabels.add(labelBuilder.build());
   }
 
   @Override
@@ -298,14 +303,13 @@ public void deleteLabel(
   }
 
   private boolean isDuplicateKey(String key) {
+    return getLabelsMap().containsKey(key);
+  }
+
+  private Map<String, Label> getLabelsMap() {
     RequestContext requestContext = RequestContext.CURRENT.get();
-    List<Label> labelList =
-        labelStore.getAllObjects(requestContext).stream()
-            .map(ContextualConfigObject::getData)
-            .collect(Collectors.toUnmodifiableList());
-    return labelList.stream()
-        .map(Label::getData)
-        .map(LabelData::getKey)
-        .anyMatch(labelKey -> labelKey.equals(key));
+    return labelStore.getAllObjects(requestContext).stream()
+        .map(ContextualConfigObject::getData)
+        .collect(Collectors.toUnmodifiableMap(label -> label.getData().getKey(), identity()));
   }
 }
diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index a6fa65eb..14bd0085 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -19,14 +19,14 @@
 import java.util.stream.Stream;
 import org.hypertrace.config.service.change.event.api.ConfigChangeEventGenerator;
 import org.hypertrace.config.service.test.MockGenericConfigService;
-import org.hypertrace.label.config.service.v1.BulkCreateLabelsRequest;
-import org.hypertrace.label.config.service.v1.BulkCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.CreateLabelRequest;
 import org.hypertrace.label.config.service.v1.CreateLabelResponse;
 import org.hypertrace.label.config.service.v1.DeleteLabelRequest;
 import org.hypertrace.label.config.service.v1.GetLabelRequest;
 import org.hypertrace.label.config.service.v1.GetLabelResponse;
 import org.hypertrace.label.config.service.v1.GetLabelsRequest;
+import org.hypertrace.label.config.service.v1.GetOrCreateLabelsRequest;
+import org.hypertrace.label.config.service.v1.GetOrCreateLabelsResponse;
 import org.hypertrace.label.config.service.v1.Label;
 import org.hypertrace.label.config.service.v1.LabelData;
 import org.hypertrace.label.config.service.v1.LabelsConfigServiceGrpc;
@@ -134,44 +134,64 @@ void test_system_createLabel() {
   }
 
   @Test
-  void test_bulkCreateLabels() {
-    List<BulkCreateLabelsRequest.LabelRequest> requests =
+  void test_getOrCreateLabels() {
+    List<GetOrCreateLabelsRequest.LabelRequest> requests =
         createLabelDataList.stream()
-            .map(data -> BulkCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
+            .map(data -> GetOrCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
             .collect(Collectors.toList());
-    BulkCreateLabelsResponse bulkCreateLabelsResponse =
-        labelConfigStub.bulkCreateLabels(
-            BulkCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
+    GetOrCreateLabelsResponse response =
+        labelConfigStub.getOrCreateLabels(
+            GetOrCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
     List<LabelData> createdLabelsList =
-        bulkCreateLabelsResponse.getLabelsList().stream()
+        response.getLabelsList().stream()
             .map(label -> label.getData())
             .collect(Collectors.toList());
     assertEquals(createLabelDataList, createdLabelsList);
   }
 
   @Test
-  void test_bulkCreateLabelsWithSystemLabels() {
-    List<BulkCreateLabelsRequest.LabelRequest> requests = new ArrayList<>();
+  void test_getOrCreateLabelsWithExistingAndSystemLabels() {
+    List<Label> initiallyCreatedLabels = createLabels();
+    List<GetOrCreateLabelsRequest.LabelRequest> requests = new ArrayList<>();
     requests.addAll(
         createLabelDataList.stream()
-            .map(data -> BulkCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
+            .map(data -> GetOrCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
             .collect(Collectors.toList()));
     requests.addAll(
         systemLabels.stream()
             .map(
                 systemLabel ->
-                    BulkCreateLabelsRequest.LabelRequest.newBuilder()
+                    GetOrCreateLabelsRequest.LabelRequest.newBuilder()
                         .setData(systemLabel.getData())
                         .build())
             .collect(Collectors.toList()));
-    BulkCreateLabelsResponse bulkCreateLabelsResponse =
-        labelConfigStub.bulkCreateLabels(
-            BulkCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
-    List<LabelData> createdLabelsList =
-        bulkCreateLabelsResponse.getLabelsList().stream()
-            .map(label -> label.getData())
-            .collect(Collectors.toList());
-    assertEquals(createLabelDataList, createdLabelsList);
+    requests.addAll(
+        Stream.of(5, 6, 7)
+            .map(
+                id ->
+                    GetOrCreateLabelsRequest.LabelRequest.newBuilder()
+                        .setData(LabelData.newBuilder().setKey("Label-" + id))
+                        .build())
+            .collect(Collectors.toList()));
+    GetOrCreateLabelsResponse response =
+        labelConfigStub.getOrCreateLabels(
+            GetOrCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
+    List<Label> expectedLabelsList = new ArrayList<>();
+    expectedLabelsList.addAll(initiallyCreatedLabels);
+    expectedLabelsList.addAll(systemLabels);
+    List<Label> actualLabelsList = response.getLabelsList();
+    assertEquals(3, actualLabelsList.size() - expectedLabelsList.size());
+    int index = 0;
+    while (index < expectedLabelsList.size()) {
+      assertEquals(expectedLabelsList.get(index), actualLabelsList.get(index));
+      index++;
+    }
+    int newLabelId = 5;
+    while (index < actualLabelsList.size()) {
+      assertEquals("Label-" + newLabelId, actualLabelsList.get(index).getData().getKey());
+      index++;
+      newLabelId++;
+    }
   }
 
   @Test

From 9fada451b81852d75cb31a7f44ad7bc2f52ad59c Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Wed, 24 Nov 2021 17:23:56 +0530
Subject: [PATCH 3/3] Address review comments

---
 .../service/LabelsConfigServiceImpl.java      | 73 ++++++++++---------
 .../service/LabelsConfigServiceImplTest.java  | 37 ++++------
 2 files changed, 52 insertions(+), 58 deletions(-)

diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index bf58889b..c24f49ea 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -13,6 +13,7 @@
 import io.grpc.stub.StreamObserver;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -109,8 +110,7 @@ public void createLabel(
       if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
         try {
           LabelData labelData = request.getData();
-          if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())
-              || isDuplicateKey(labelData.getKey())) {
+          if (isDuplicateKey(requestContext, labelData.getKey())) {
             // Creating a label with a name that clashes with one of system labels name
             responseObserver.onError(new StatusRuntimeException(Status.ALREADY_EXISTS));
             return;
@@ -146,21 +146,33 @@ public void getOrCreateLabels(
       Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
       if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
         try {
-          final Map<String, Label> existingLabelsMap = getLabelsMap();
-          final List<Label> newLabels = new ArrayList<>();
-          final List<Label> allLabels = new ArrayList<>();
-          request.getRequestsList().stream()
-              .forEach(
-                  labelRequest ->
-                      this.buildLabelFromRequest(
-                          labelRequest, existingLabelsMap, newLabels, allLabels));
+          final Map<String, Label> existingLabelsMap = getLabelsMap(requestContext);
+          List<Label> newLabels =
+              request.getRequestsList().stream()
+                  .filter(
+                      labelRequest ->
+                          !existingLabelsMap.containsKey(labelRequest.getData().getKey()))
+                  .map(this::buildLabelFromRequest)
+                  .collect(Collectors.toList());
+          Map<String, Label> createdLabelsMap;
           if (!newLabels.isEmpty()) {
-            List<Label> createdLabels =
+            createdLabelsMap =
                 labelStore.upsertObjects(requestContext, newLabels).stream()
                     .map(org.hypertrace.config.objectstore.ConfigObject::getData)
-                    .collect(Collectors.toList());
-            allLabels.addAll(createdLabels);
+                    .collect(
+                        Collectors.toUnmodifiableMap(
+                            label -> label.getData().getKey(), identity()));
+          } else {
+            createdLabelsMap = Collections.emptyMap();
           }
+          final Map<String, Label> allLabelsMap = new HashMap<>();
+          allLabelsMap.putAll(existingLabelsMap);
+          allLabelsMap.putAll(createdLabelsMap);
+          List<Label> allLabels =
+              request.getRequestsList().stream()
+                  .map(GetOrCreateLabelsRequest.LabelRequest::getData)
+                  .map(data -> allLabelsMap.get(data.getKey()))
+                  .collect(Collectors.toList());
           responseObserver.onNext(
               GetOrCreateLabelsResponse.newBuilder().addAllLabels(allLabels).build());
           responseObserver.onCompleted();
@@ -175,26 +187,14 @@ public void getOrCreateLabels(
     }
   }
 
-  private void buildLabelFromRequest(
-      GetOrCreateLabelsRequest.LabelRequest request,
-      Map<String, Label> existingLabelsMap,
-      List<Label> newLabels,
-      List<Label> allLabels) {
+  private Label buildLabelFromRequest(GetOrCreateLabelsRequest.LabelRequest request) {
     LabelData labelData = request.getData();
-    if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())) {
-      allLabels.add(systemLabelsKeyLabelMap.get(labelData.getKey()));
-      return;
-    }
-    if (existingLabelsMap.containsKey(labelData.getKey())) {
-      allLabels.add(existingLabelsMap.get(labelData.getKey()));
-      return;
-    }
     Label.Builder labelBuilder =
         Label.newBuilder().setId(UUID.randomUUID().toString()).setData(labelData);
     if (request.hasCreatedByApplicationRuleId()) {
       labelBuilder.setCreatedByApplicationRuleId(request.getCreatedByApplicationRuleId());
     }
-    newLabels.add(labelBuilder.build());
+    return labelBuilder.build();
   }
 
   @Override
@@ -241,13 +241,12 @@ public void updateLabel(
       Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
       if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
         try {
-          if (systemLabelsIdLabelMap.containsKey(request.getId())
-              || systemLabelsKeyLabelMap.containsKey(updateLabelData.getKey())) {
+          if (systemLabelsIdLabelMap.containsKey(request.getId())) {
             // Updating a system label will error
             responseObserver.onError(new StatusRuntimeException(Status.INVALID_ARGUMENT));
             return;
           }
-          if (isDuplicateKey(updateLabelData.getKey())) {
+          if (isDuplicateKey(requestContext, updateLabelData.getKey())) {
             responseObserver.onError(new StatusRuntimeException(Status.ALREADY_EXISTS));
             return;
           }
@@ -302,14 +301,16 @@ public void deleteLabel(
     }
   }
 
-  private boolean isDuplicateKey(String key) {
-    return getLabelsMap().containsKey(key);
+  private boolean isDuplicateKey(RequestContext requestContext, String key) {
+    return getLabelsMap(requestContext).containsKey(key);
   }
 
-  private Map<String, Label> getLabelsMap() {
-    RequestContext requestContext = RequestContext.CURRENT.get();
-    return labelStore.getAllObjects(requestContext).stream()
+  private Map<String, Label> getLabelsMap(RequestContext requestContext) {
+    Map<String, Label> existingLabelsMap = new HashMap<>();
+    existingLabelsMap.putAll(systemLabelsKeyLabelMap);
+    labelStore.getAllObjects(requestContext).stream()
         .map(ContextualConfigObject::getData)
-        .collect(Collectors.toUnmodifiableMap(label -> label.getData().getKey(), identity()));
+        .forEach(label -> existingLabelsMap.put(label.getData().getKey(), label));
+    return Collections.unmodifiableMap(existingLabelsMap);
   }
 }
diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index 14bd0085..85f1ab6f 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -153,6 +153,14 @@ void test_getOrCreateLabels() {
   void test_getOrCreateLabelsWithExistingAndSystemLabels() {
     List<Label> initiallyCreatedLabels = createLabels();
     List<GetOrCreateLabelsRequest.LabelRequest> requests = new ArrayList<>();
+    requests.addAll(
+        Stream.of(5, 6, 7)
+            .map(
+                id ->
+                    GetOrCreateLabelsRequest.LabelRequest.newBuilder()
+                        .setData(LabelData.newBuilder().setKey("Label-" + id))
+                        .build())
+            .collect(Collectors.toList()));
     requests.addAll(
         createLabelDataList.stream()
             .map(data -> GetOrCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
@@ -165,33 +173,18 @@ void test_getOrCreateLabelsWithExistingAndSystemLabels() {
                         .setData(systemLabel.getData())
                         .build())
             .collect(Collectors.toList()));
-    requests.addAll(
-        Stream.of(5, 6, 7)
-            .map(
-                id ->
-                    GetOrCreateLabelsRequest.LabelRequest.newBuilder()
-                        .setData(LabelData.newBuilder().setKey("Label-" + id))
-                        .build())
-            .collect(Collectors.toList()));
     GetOrCreateLabelsResponse response =
         labelConfigStub.getOrCreateLabels(
             GetOrCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
+    List<Label> actualLabelsList = new ArrayList<>(response.getLabelsList());
+    Stream.of(5, 6, 7)
+        .forEach(
+            newLabelId ->
+                assertEquals("Label-" + newLabelId, actualLabelsList.remove(0).getData().getKey()));
     List<Label> expectedLabelsList = new ArrayList<>();
     expectedLabelsList.addAll(initiallyCreatedLabels);
     expectedLabelsList.addAll(systemLabels);
-    List<Label> actualLabelsList = response.getLabelsList();
-    assertEquals(3, actualLabelsList.size() - expectedLabelsList.size());
-    int index = 0;
-    while (index < expectedLabelsList.size()) {
-      assertEquals(expectedLabelsList.get(index), actualLabelsList.get(index));
-      index++;
-    }
-    int newLabelId = 5;
-    while (index < actualLabelsList.size()) {
-      assertEquals("Label-" + newLabelId, actualLabelsList.get(index).getData().getKey());
-      index++;
-      newLabelId++;
-    }
+    assertEquals(expectedLabelsList, actualLabelsList);
   }
 
   @Test
@@ -328,7 +321,7 @@ void test_system_updateLabel() {
               () -> {
                 labelConfigStub.updateLabel(updateLabelRequest);
               });
-      assertEquals(Status.INVALID_ARGUMENT, Status.fromThrowable(exception));
+      assertEquals(Status.ALREADY_EXISTS, Status.fromThrowable(exception));
     }
   }