Skip to content

Commit d92d6eb

Browse files
saxenakshitizsaxenakshitiz
authored
Changes to support bulk create of labels (#82)
1 parent f8ee358 commit d92d6eb

File tree

3 files changed

+161
-22
lines changed

3 files changed

+161
-22
lines changed

‎labels-config-service-api/src/main/proto/org/hypertrace/label/config/service/v1/labels_config_service.proto‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ service LabelsConfigService {
1111
// create label
1212
rpc CreateLabel(CreateLabelRequest) returns (CreateLabelResponse) {}
1313

14+
// get or create labels
15+
rpc GetOrCreateLabels(GetOrCreateLabelsRequest) returns (GetOrCreateLabelsResponse) {}
16+
1417
// get label by id
1518
rpc GetLabel(GetLabelRequest) returns (GetLabelResponse) {}
1619

@@ -34,6 +37,19 @@ message CreateLabelResponse {
3437
Label label = 1;
3538
}
3639

40+
message GetOrCreateLabelsRequest {
41+
repeated LabelRequest requests = 1;
42+
43+
message LabelRequest {
44+
LabelData data = 1;
45+
optional string created_by_application_rule_id = 2;
46+
}
47+
}
48+
49+
message GetOrCreateLabelsResponse {
50+
repeated Label labels = 1;
51+
}
52+
3753
message GetLabelRequest {
3854
string id = 1;
3955
}

‎labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java‎

Lines changed: 81 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package org.hypertrace.label.config.service;
22

3+
import static java.util.function.Function.identity;
4+
35
import com.google.common.util.concurrent.Striped;
46
import com.google.protobuf.Duration;
57
import com.google.protobuf.util.JsonFormat;
@@ -11,12 +13,12 @@
1113
import io.grpc.stub.StreamObserver;
1214
import java.util.ArrayList;
1315
import java.util.Collections;
16+
import java.util.HashMap;
1417
import java.util.List;
1518
import java.util.Map;
1619
import java.util.UUID;
1720
import java.util.concurrent.TimeUnit;
1821
import java.util.concurrent.locks.Lock;
19-
import java.util.function.Function;
2022
import java.util.stream.Collectors;
2123
import lombok.SneakyThrows;
2224
import lombok.extern.slf4j.Slf4j;
@@ -33,6 +35,8 @@
3335
import org.hypertrace.label.config.service.v1.GetLabelResponse;
3436
import org.hypertrace.label.config.service.v1.GetLabelsRequest;
3537
import org.hypertrace.label.config.service.v1.GetLabelsResponse;
38+
import org.hypertrace.label.config.service.v1.GetOrCreateLabelsRequest;
39+
import org.hypertrace.label.config.service.v1.GetOrCreateLabelsResponse;
3640
import org.hypertrace.label.config.service.v1.Label;
3741
import org.hypertrace.label.config.service.v1.LabelData;
3842
import org.hypertrace.label.config.service.v1.LabelsConfigServiceGrpc;
@@ -71,13 +75,11 @@ public LabelsConfigServiceImpl(
7175
if (systemLabelsObjectList != null) {
7276
systemLabels = buildSystemLabelList(systemLabelsObjectList);
7377
systemLabelsIdLabelMap =
74-
systemLabels.stream()
75-
.collect(Collectors.toUnmodifiableMap(Label::getId, Function.identity()));
78+
systemLabels.stream().collect(Collectors.toUnmodifiableMap(Label::getId, identity()));
7679
systemLabelsKeyLabelMap =
7780
systemLabels.stream()
7881
.collect(
79-
Collectors.toUnmodifiableMap(
80-
(label) -> label.getData().getKey(), Function.identity()));
82+
Collectors.toUnmodifiableMap((label) -> label.getData().getKey(), identity()));
8183
} else {
8284
systemLabels = Collections.emptyList();
8385
systemLabelsIdLabelMap = Collections.emptyMap();
@@ -108,8 +110,7 @@ public void createLabel(
108110
if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
109111
try {
110112
LabelData labelData = request.getData();
111-
if (systemLabelsKeyLabelMap.containsKey(labelData.getKey())
112-
|| isDuplicateKey(labelData.getKey())) {
113+
if (isDuplicateKey(requestContext, labelData.getKey())) {
113114
// Creating a label with a name that clashes with one of system labels name
114115
responseObserver.onError(new StatusRuntimeException(Status.ALREADY_EXISTS));
115116
return;
@@ -136,6 +137,66 @@ public void createLabel(
136137
}
137138
}
138139

140+
@Override
141+
public void getOrCreateLabels(
142+
GetOrCreateLabelsRequest request,
143+
StreamObserver<GetOrCreateLabelsResponse> responseObserver) {
144+
try {
145+
RequestContext requestContext = RequestContext.CURRENT.get();
146+
Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
147+
if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
148+
try {
149+
final Map<String, Label> existingLabelsMap = getLabelsMap(requestContext);
150+
List<Label> newLabels =
151+
request.getRequestsList().stream()
152+
.filter(
153+
labelRequest ->
154+
!existingLabelsMap.containsKey(labelRequest.getData().getKey()))
155+
.map(this::buildLabelFromRequest)
156+
.collect(Collectors.toList());
157+
Map<String, Label> createdLabelsMap;
158+
if (!newLabels.isEmpty()) {
159+
createdLabelsMap =
160+
labelStore.upsertObjects(requestContext, newLabels).stream()
161+
.map(org.hypertrace.config.objectstore.ConfigObject::getData)
162+
.collect(
163+
Collectors.toUnmodifiableMap(
164+
label -> label.getData().getKey(), identity()));
165+
} else {
166+
createdLabelsMap = Collections.emptyMap();
167+
}
168+
final Map<String, Label> allLabelsMap = new HashMap<>();
169+
allLabelsMap.putAll(existingLabelsMap);
170+
allLabelsMap.putAll(createdLabelsMap);
171+
List<Label> allLabels =
172+
request.getRequestsList().stream()
173+
.map(GetOrCreateLabelsRequest.LabelRequest::getData)
174+
.map(data -> allLabelsMap.get(data.getKey()))
175+
.collect(Collectors.toList());
176+
responseObserver.onNext(
177+
GetOrCreateLabelsResponse.newBuilder().addAllLabels(allLabels).build());
178+
responseObserver.onCompleted();
179+
} finally {
180+
labelsLock.unlock();
181+
}
182+
} else {
183+
responseObserver.onError(new StatusRuntimeException(Status.ABORTED));
184+
}
185+
} catch (Exception e) {
186+
responseObserver.onError(e);
187+
}
188+
}
189+
190+
private Label buildLabelFromRequest(GetOrCreateLabelsRequest.LabelRequest request) {
191+
LabelData labelData = request.getData();
192+
Label.Builder labelBuilder =
193+
Label.newBuilder().setId(UUID.randomUUID().toString()).setData(labelData);
194+
if (request.hasCreatedByApplicationRuleId()) {
195+
labelBuilder.setCreatedByApplicationRuleId(request.getCreatedByApplicationRuleId());
196+
}
197+
return labelBuilder.build();
198+
}
199+
139200
@Override
140201
public void getLabel(GetLabelRequest request, StreamObserver<GetLabelResponse> responseObserver) {
141202
RequestContext requestContext = RequestContext.CURRENT.get();
@@ -180,13 +241,12 @@ public void updateLabel(
180241
Lock labelsLock = this.stripedLabelsLock.get(requestContext.getTenantId());
181242
if (labelsLock.tryLock(WAIT_TIME.getSeconds(), TimeUnit.SECONDS)) {
182243
try {
183-
if (systemLabelsIdLabelMap.containsKey(request.getId())
184-
|| systemLabelsKeyLabelMap.containsKey(updateLabelData.getKey())) {
244+
if (systemLabelsIdLabelMap.containsKey(request.getId())) {
185245
// Updating a system label will error
186246
responseObserver.onError(new StatusRuntimeException(Status.INVALID_ARGUMENT));
187247
return;
188248
}
189-
if (isDuplicateKey(updateLabelData.getKey())) {
249+
if (isDuplicateKey(requestContext, updateLabelData.getKey())) {
190250
responseObserver.onError(new StatusRuntimeException(Status.ALREADY_EXISTS));
191251
return;
192252
}
@@ -241,15 +301,16 @@ public void deleteLabel(
241301
}
242302
}
243303

244-
private boolean isDuplicateKey(String key) {
245-
RequestContext requestContext = RequestContext.CURRENT.get();
246-
List<Label> labelList =
247-
labelStore.getAllObjects(requestContext).stream()
248-
.map(ContextualConfigObject::getData)
249-
.collect(Collectors.toUnmodifiableList());
250-
return labelList.stream()
251-
.map(Label::getData)
252-
.map(LabelData::getKey)
253-
.anyMatch(labelKey -> labelKey.equals(key));
304+
private boolean isDuplicateKey(RequestContext requestContext, String key) {
305+
return getLabelsMap(requestContext).containsKey(key);
306+
}
307+
308+
private Map<String, Label> getLabelsMap(RequestContext requestContext) {
309+
Map<String, Label> existingLabelsMap = new HashMap<>();
310+
existingLabelsMap.putAll(systemLabelsKeyLabelMap);
311+
labelStore.getAllObjects(requestContext).stream()
312+
.map(ContextualConfigObject::getData)
313+
.forEach(label -> existingLabelsMap.put(label.getData().getKey(), label));
314+
return Collections.unmodifiableMap(existingLabelsMap);
254315
}
255316
}

‎labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java‎

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import io.grpc.Channel;
1010
import io.grpc.Status;
1111
import io.grpc.StatusRuntimeException;
12+
import java.util.ArrayList;
1213
import java.util.Arrays;
1314
import java.util.HashMap;
1415
import java.util.List;
@@ -24,6 +25,8 @@
2425
import org.hypertrace.label.config.service.v1.GetLabelRequest;
2526
import org.hypertrace.label.config.service.v1.GetLabelResponse;
2627
import org.hypertrace.label.config.service.v1.GetLabelsRequest;
28+
import org.hypertrace.label.config.service.v1.GetOrCreateLabelsRequest;
29+
import org.hypertrace.label.config.service.v1.GetOrCreateLabelsResponse;
2730
import org.hypertrace.label.config.service.v1.Label;
2831
import org.hypertrace.label.config.service.v1.LabelData;
2932
import org.hypertrace.label.config.service.v1.LabelsConfigServiceGrpc;
@@ -55,7 +58,12 @@ public final class LabelsConfigServiceImplTest {
5558
@BeforeEach
5659
void setUp() {
5760
mockGenericConfigService =
58-
new MockGenericConfigService().mockUpsert().mockGet().mockGetAll().mockDelete();
61+
new MockGenericConfigService()
62+
.mockUpsert()
63+
.mockUpsertAll()
64+
.mockGet()
65+
.mockGetAll()
66+
.mockDelete();
5967
Map<String, List<Map<String, String>>> systemLabelsConfigMap = new HashMap<>();
6068
systemLabelsConfigMap.put(
6169
"system.labels",
@@ -125,6 +133,60 @@ void test_system_createLabel() {
125133
}
126134
}
127135

136+
@Test
137+
void test_getOrCreateLabels() {
138+
List<GetOrCreateLabelsRequest.LabelRequest> requests =
139+
createLabelDataList.stream()
140+
.map(data -> GetOrCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
141+
.collect(Collectors.toList());
142+
GetOrCreateLabelsResponse response =
143+
labelConfigStub.getOrCreateLabels(
144+
GetOrCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
145+
List<LabelData> createdLabelsList =
146+
response.getLabelsList().stream()
147+
.map(label -> label.getData())
148+
.collect(Collectors.toList());
149+
assertEquals(createLabelDataList, createdLabelsList);
150+
}
151+
152+
@Test
153+
void test_getOrCreateLabelsWithExistingAndSystemLabels() {
154+
List<Label> initiallyCreatedLabels = createLabels();
155+
List<GetOrCreateLabelsRequest.LabelRequest> requests = new ArrayList<>();
156+
requests.addAll(
157+
Stream.of(5, 6, 7)
158+
.map(
159+
id ->
160+
GetOrCreateLabelsRequest.LabelRequest.newBuilder()
161+
.setData(LabelData.newBuilder().setKey("Label-" + id))
162+
.build())
163+
.collect(Collectors.toList()));
164+
requests.addAll(
165+
createLabelDataList.stream()
166+
.map(data -> GetOrCreateLabelsRequest.LabelRequest.newBuilder().setData(data).build())
167+
.collect(Collectors.toList()));
168+
requests.addAll(
169+
systemLabels.stream()
170+
.map(
171+
systemLabel ->
172+
GetOrCreateLabelsRequest.LabelRequest.newBuilder()
173+
.setData(systemLabel.getData())
174+
.build())
175+
.collect(Collectors.toList()));
176+
GetOrCreateLabelsResponse response =
177+
labelConfigStub.getOrCreateLabels(
178+
GetOrCreateLabelsRequest.newBuilder().addAllRequests(requests).build());
179+
List<Label> actualLabelsList = new ArrayList<>(response.getLabelsList());
180+
Stream.of(5, 6, 7)
181+
.forEach(
182+
newLabelId ->
183+
assertEquals("Label-" + newLabelId, actualLabelsList.remove(0).getData().getKey()));
184+
List<Label> expectedLabelsList = new ArrayList<>();
185+
expectedLabelsList.addAll(initiallyCreatedLabels);
186+
expectedLabelsList.addAll(systemLabels);
187+
assertEquals(expectedLabelsList, actualLabelsList);
188+
}
189+
128190
@Test
129191
void test_getLabel() {
130192
List<Label> createdLabelsList =
@@ -259,7 +321,7 @@ void test_system_updateLabel() {
259321
() -> {
260322
labelConfigStub.updateLabel(updateLabelRequest);
261323
});
262-
assertEquals(Status.INVALID_ARGUMENT, Status.fromThrowable(exception));
324+
assertEquals(Status.ALREADY_EXISTS, Status.fromThrowable(exception));
263325
}
264326
}
265327

0 commit comments

Comments
 (0)