Skip to content

Commit e7559b7

Browse files
pthirunpthirun
committed
[Controller] Add gRPC support for getClusterHealthStores API
Add gRPC support for the getClusterHealthStores endpoint while maintaining backward compatibility with the existing HTTP endpoint. Changes: - Add getClusterHealthStores RPC to StoreGrpcService.proto - Add handler method to StoreRequestHandler - Add gRPC service implementation - Update HTTP route to convert requests to gRPC format - Add comprehensive unit and integration tests
1 parent 406454c commit e7559b7

8 files changed

Lines changed: 332 additions & 3 deletions

File tree

‎internal/venice-common/src/main/proto/controller/StoreGrpcService.proto‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ service StoreGrpcService {
1414
rpc checkResourceCleanupForStoreCreation(ClusterStoreGrpcInfo) returns (ResourceCleanupCheckGrpcResponse) {}
1515
rpc validateStoreDeleted(ValidateStoreDeletedGrpcRequest) returns (ValidateStoreDeletedGrpcResponse);
1616
rpc listStores(ListStoresGrpcRequest) returns (ListStoresGrpcResponse);
17+
rpc getClusterHealthStores(GetClusterHealthStoresGrpcRequest) returns (GetClusterHealthStoresGrpcResponse);
1718
}
1819

1920
message CreateStoreGrpcRequest {
@@ -82,4 +83,13 @@ message ListStoresGrpcRequest {
8283
message ListStoresGrpcResponse {
8384
string clusterName = 1;
8485
repeated string storeNames = 2;
86+
}
87+
88+
message GetClusterHealthStoresGrpcRequest {
89+
string clusterName = 1;
90+
}
91+
92+
message GetClusterHealthStoresGrpcResponse {
93+
string clusterName = 1;
94+
map<string, string> storeStatusMap = 2;
8595
}

‎internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestControllerGrpcEndpoints.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,13 @@
2020
import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse;
2121
import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest;
2222
import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse;
23+
<<<<<<< HEAD
2324
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
2425
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
26+
=======
27+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
28+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
29+
>>>>>>> 3b44a8ea8 ([Controller] Add gRPC support for getClusterHealthStores API)
2530
import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest;
2631
import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse;
2732
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
@@ -386,6 +391,66 @@ public void testGetValueSchemaGrpcEndpoint() {
386391
assertEquals(exception.getStatus().getCode(), io.grpc.Status.Code.INVALID_ARGUMENT);
387392
}
388393

394+
@Test(timeOut = TIMEOUT_MS)
395+
public void testGetClusterHealthStoresGrpcEndpoint() {
396+
String storeName1 = Utils.getUniqueString("test_health_stores_1");
397+
String storeName2 = Utils.getUniqueString("test_health_stores_2");
398+
String controllerGrpcUrl = veniceCluster.getLeaderVeniceController().getControllerGrpcUrl();
399+
ManagedChannel channel = Grpc.newChannelBuilder(controllerGrpcUrl, InsecureChannelCredentials.create()).build();
400+
StoreGrpcServiceGrpc.StoreGrpcServiceBlockingStub storeBlockingStub = StoreGrpcServiceGrpc.newBlockingStub(channel);
401+
402+
// Step 1: Create two stores
403+
CreateStoreGrpcRequest createStoreRequest1 = CreateStoreGrpcRequest.newBuilder()
404+
.setStoreInfo(
405+
ClusterStoreGrpcInfo.newBuilder()
406+
.setClusterName(veniceCluster.getClusterName())
407+
.setStoreName(storeName1)
408+
.build())
409+
.setOwner("owner")
410+
.setKeySchema(DEFAULT_KEY_SCHEMA)
411+
.setValueSchema("\"string\"")
412+
.build();
413+
CreateStoreGrpcResponse createResponse1 = storeBlockingStub.createStore(createStoreRequest1);
414+
assertNotNull(createResponse1, "Response should not be null");
415+
416+
CreateStoreGrpcRequest createStoreRequest2 = CreateStoreGrpcRequest.newBuilder()
417+
.setStoreInfo(
418+
ClusterStoreGrpcInfo.newBuilder()
419+
.setClusterName(veniceCluster.getClusterName())
420+
.setStoreName(storeName2)
421+
.build())
422+
.setOwner("owner")
423+
.setKeySchema(DEFAULT_KEY_SCHEMA)
424+
.setValueSchema("\"string\"")
425+
.build();
426+
CreateStoreGrpcResponse createResponse2 = storeBlockingStub.createStore(createStoreRequest2);
427+
assertNotNull(createResponse2, "Response should not be null");
428+
429+
// Step 2: Get cluster health stores
430+
GetClusterHealthStoresGrpcRequest healthRequest =
431+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName(veniceCluster.getClusterName()).build();
432+
433+
GetClusterHealthStoresGrpcResponse healthResponse = storeBlockingStub.getClusterHealthStores(healthRequest);
434+
assertNotNull(healthResponse, "Response should not be null");
435+
assertEquals(healthResponse.getClusterName(), veniceCluster.getClusterName());
436+
437+
// Verify the stores we created are in the status map
438+
assertTrue(
439+
healthResponse.getStoreStatusMapMap().containsKey(storeName1),
440+
"Store status map should contain " + storeName1);
441+
assertTrue(
442+
healthResponse.getStoreStatusMapMap().containsKey(storeName2),
443+
"Store status map should contain " + storeName2);
444+
445+
// Verify the statuses are not null/empty
446+
assertNotNull(
447+
healthResponse.getStoreStatusMapMap().get(storeName1),
448+
"Status for " + storeName1 + " should not be null");
449+
assertNotNull(
450+
healthResponse.getStoreStatusMapMap().get(storeName2),
451+
"Status for " + storeName2 + " should not be null");
452+
}
453+
389454
private static class MockDynamicAccessController extends NoOpDynamicAccessController {
390455
private final Set<String> resourcesInAllowList = ConcurrentHashMap.newKeySet();
391456

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/grpc/server/StoreGrpcServiceImpl.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
import com.linkedin.venice.protocols.controller.DeleteAclForStoreGrpcResponse;
1515
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcRequest;
1616
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcResponse;
17+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
18+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
1719
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
1820
import com.linkedin.venice.protocols.controller.ListStoresGrpcResponse;
1921
import com.linkedin.venice.protocols.controller.ResourceCleanupCheckGrpcResponse;
@@ -147,4 +149,22 @@ public void listStores(ListStoresGrpcRequest grpcRequest, StreamObserver<ListSto
147149
clusterName,
148150
null);
149151
}
152+
153+
/**
154+
* Gets the health status of all stores in the specified cluster.
155+
* No ACL check; any user can query store health statuses.
156+
*/
157+
@Override
158+
public void getClusterHealthStores(
159+
GetClusterHealthStoresGrpcRequest grpcRequest,
160+
StreamObserver<GetClusterHealthStoresGrpcResponse> responseObserver) {
161+
LOGGER.debug("Received getClusterHealthStores with args: {}", grpcRequest);
162+
String clusterName = grpcRequest.getClusterName();
163+
handleRequest(
164+
StoreGrpcServiceGrpc.getGetClusterHealthStoresMethod(),
165+
() -> storeRequestHandler.getClusterHealthStores(grpcRequest),
166+
responseObserver,
167+
clusterName,
168+
null);
169+
}
150170
}

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/server/StoreRequestHandler.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
import com.linkedin.venice.protocols.controller.DeleteAclForStoreGrpcResponse;
1414
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcRequest;
1515
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcResponse;
16+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
17+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
1618
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
1719
import com.linkedin.venice.protocols.controller.ListStoresGrpcResponse;
1820
import com.linkedin.venice.protocols.controller.UpdateAclForStoreGrpcRequest;
@@ -22,6 +24,7 @@
2224
import com.linkedin.venice.systemstore.schemas.StoreProperties;
2325
import java.util.ArrayList;
2426
import java.util.List;
27+
import java.util.Map;
2528
import java.util.Optional;
2629
import org.apache.avro.Schema;
2730
import org.apache.commons.lang.StringUtils;
@@ -255,4 +258,25 @@ public ListStoresGrpcResponse listStores(ListStoresGrpcRequest request) {
255258
LOGGER.info("Found {} stores in cluster: {}", selectedStoreNames.size(), clusterName);
256259
return ListStoresGrpcResponse.newBuilder().setClusterName(clusterName).addAllStoreNames(selectedStoreNames).build();
257260
}
261+
262+
/**
263+
* Gets the health status of all stores in the specified cluster.
264+
* @param request the request containing cluster name
265+
* @return response containing the map of store names to their statuses
266+
*/
267+
public GetClusterHealthStoresGrpcResponse getClusterHealthStores(GetClusterHealthStoresGrpcRequest request) {
268+
String clusterName = request.getClusterName();
269+
if (StringUtils.isBlank(clusterName)) {
270+
throw new IllegalArgumentException("Cluster name is required");
271+
}
272+
273+
LOGGER.info("Getting health status for all stores in cluster: {}", clusterName);
274+
Map<String, String> storeStatusMap = admin.getAllStoreStatuses(clusterName);
275+
LOGGER.info("Found {} stores with health status in cluster: {}", storeStatusMap.size(), clusterName);
276+
277+
return GetClusterHealthStoresGrpcResponse.newBuilder()
278+
.setClusterName(clusterName)
279+
.putAllStoreStatusMap(storeStatusMap)
280+
.build();
281+
}
258282
}

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/server/StoresRoutes.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,8 @@
109109
import com.linkedin.venice.meta.StoreInfo;
110110
import com.linkedin.venice.meta.Version;
111111
import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo;
112+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
113+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
112114
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
113115
import com.linkedin.venice.protocols.controller.ListStoresGrpcResponse;
114116
import com.linkedin.venice.protocols.controller.ValidateStoreDeletedGrpcRequest;
@@ -228,9 +230,13 @@ public Route getAllStoresStatuses(Admin admin) {
228230
public void internalHandle(Request request, MultiStoreStatusResponse veniceResponse) {
229231
AdminSparkServer.validateParams(request, CLUSTER_HEALTH_STORES.getParams(), admin);
230232
String clusterName = request.queryParams(CLUSTER);
231-
veniceResponse.setCluster(clusterName);
232-
Map<String, String> storeStatusMap = admin.getAllStoreStatuses(clusterName);
233-
veniceResponse.setStoreStatusMap(storeStatusMap);
233+
234+
GetClusterHealthStoresGrpcRequest grpcRequest =
235+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName(clusterName).build();
236+
GetClusterHealthStoresGrpcResponse grpcResponse = storeRequestHandler.getClusterHealthStores(grpcRequest);
237+
238+
veniceResponse.setCluster(grpcResponse.getClusterName());
239+
veniceResponse.setStoreStatusMap(grpcResponse.getStoreStatusMapMap());
234240
}
235241
};
236242
}

‎services/venice-controller/src/test/java/com/linkedin/venice/controller/grpc/server/StoreGrpcServiceImplTest.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import com.linkedin.venice.protocols.controller.DeleteAclForStoreGrpcResponse;
2626
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcRequest;
2727
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcResponse;
28+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
29+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
2830
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
2931
import com.linkedin.venice.protocols.controller.ListStoresGrpcResponse;
3032
import com.linkedin.venice.protocols.controller.ResourceCleanupCheckGrpcResponse;
@@ -42,6 +44,8 @@
4244
import io.grpc.inprocess.InProcessChannelBuilder;
4345
import io.grpc.inprocess.InProcessServerBuilder;
4446
import java.util.Arrays;
47+
import java.util.HashMap;
48+
import java.util.Map;
4549
import org.testng.annotations.AfterMethod;
4650
import org.testng.annotations.BeforeMethod;
4751
import org.testng.annotations.Test;
@@ -445,4 +449,80 @@ public void testListStoresWithFilters() {
445449
assertEquals(actualResponse.getClusterName(), TEST_CLUSTER, "Cluster name should match");
446450
assertEquals(actualResponse.getStoreNamesCount(), 1, "Should have 1 store after filtering");
447451
}
452+
453+
@Test
454+
public void testGetClusterHealthStoresReturnsSuccessfulResponse() {
455+
GetClusterHealthStoresGrpcRequest request =
456+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName(TEST_CLUSTER).build();
457+
Map<String, String> storeStatusMap = new HashMap<>();
458+
storeStatusMap.put("store1", "ONLINE");
459+
storeStatusMap.put("store2", "DEGRADED");
460+
storeStatusMap.put("store3", "UNAVAILABLE");
461+
GetClusterHealthStoresGrpcResponse expectedResponse = GetClusterHealthStoresGrpcResponse.newBuilder()
462+
.setClusterName(TEST_CLUSTER)
463+
.putAllStoreStatusMap(storeStatusMap)
464+
.build();
465+
when(storeRequestHandler.getClusterHealthStores(any(GetClusterHealthStoresGrpcRequest.class)))
466+
.thenReturn(expectedResponse);
467+
468+
GetClusterHealthStoresGrpcResponse actualResponse = blockingStub.getClusterHealthStores(request);
469+
470+
assertNotNull(actualResponse, "Response should not be null");
471+
assertEquals(actualResponse.getClusterName(), TEST_CLUSTER, "Cluster name should match");
472+
assertEquals(actualResponse.getStoreStatusMapCount(), 3, "Should have 3 stores");
473+
assertEquals(actualResponse.getStoreStatusMapMap().get("store1"), "ONLINE", "store1 should be ONLINE");
474+
assertEquals(actualResponse.getStoreStatusMapMap().get("store2"), "DEGRADED", "store2 should be DEGRADED");
475+
assertEquals(actualResponse.getStoreStatusMapMap().get("store3"), "UNAVAILABLE", "store3 should be UNAVAILABLE");
476+
}
477+
478+
@Test
479+
public void testGetClusterHealthStoresReturnsEmptyMapWhenNoStores() {
480+
GetClusterHealthStoresGrpcRequest request =
481+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName(TEST_CLUSTER).build();
482+
GetClusterHealthStoresGrpcResponse expectedResponse =
483+
GetClusterHealthStoresGrpcResponse.newBuilder().setClusterName(TEST_CLUSTER).build();
484+
when(storeRequestHandler.getClusterHealthStores(any(GetClusterHealthStoresGrpcRequest.class)))
485+
.thenReturn(expectedResponse);
486+
487+
GetClusterHealthStoresGrpcResponse actualResponse = blockingStub.getClusterHealthStores(request);
488+
489+
assertNotNull(actualResponse, "Response should not be null");
490+
assertEquals(actualResponse.getClusterName(), TEST_CLUSTER, "Cluster name should match");
491+
assertEquals(actualResponse.getStoreStatusMapCount(), 0, "Should have 0 stores");
492+
}
493+
494+
@Test
495+
public void testGetClusterHealthStoresReturnsErrorResponse() {
496+
GetClusterHealthStoresGrpcRequest request =
497+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName(TEST_CLUSTER).build();
498+
when(storeRequestHandler.getClusterHealthStores(any(GetClusterHealthStoresGrpcRequest.class)))
499+
.thenThrow(new VeniceException("Failed to get cluster health stores"));
500+
501+
StatusRuntimeException e =
502+
expectThrows(StatusRuntimeException.class, () -> blockingStub.getClusterHealthStores(request));
503+
504+
assertNotNull(e.getStatus(), "Status should not be null");
505+
assertEquals(e.getStatus().getCode(), Status.INTERNAL.getCode());
506+
VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e);
507+
assertNotNull(errorInfo, "Error info should not be null");
508+
assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR);
509+
assertTrue(errorInfo.getErrorMessage().contains("Failed to get cluster health stores"));
510+
}
511+
512+
@Test
513+
public void testGetClusterHealthStoresReturnsBadRequestForMissingClusterName() {
514+
GetClusterHealthStoresGrpcRequest request = GetClusterHealthStoresGrpcRequest.newBuilder().build();
515+
when(storeRequestHandler.getClusterHealthStores(any(GetClusterHealthStoresGrpcRequest.class)))
516+
.thenThrow(new IllegalArgumentException("Cluster name is required"));
517+
518+
StatusRuntimeException e =
519+
expectThrows(StatusRuntimeException.class, () -> blockingStub.getClusterHealthStores(request));
520+
521+
assertNotNull(e.getStatus(), "Status should not be null");
522+
assertEquals(e.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode());
523+
VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e);
524+
assertNotNull(errorInfo, "Error info should not be null");
525+
assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST);
526+
assertTrue(errorInfo.getErrorMessage().contains("Cluster name is required"));
527+
}
448528
}

‎services/venice-controller/src/test/java/com/linkedin/venice/controller/server/StoreRequestHandlerTest.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,17 @@
2222
import com.linkedin.venice.protocols.controller.DeleteAclForStoreGrpcResponse;
2323
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcRequest;
2424
import com.linkedin.venice.protocols.controller.GetAclForStoreGrpcResponse;
25+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcRequest;
26+
import com.linkedin.venice.protocols.controller.GetClusterHealthStoresGrpcResponse;
2527
import com.linkedin.venice.protocols.controller.ListStoresGrpcRequest;
2628
import com.linkedin.venice.protocols.controller.ListStoresGrpcResponse;
2729
import com.linkedin.venice.protocols.controller.UpdateAclForStoreGrpcRequest;
2830
import com.linkedin.venice.protocols.controller.UpdateAclForStoreGrpcResponse;
2931
import java.util.Arrays;
3032
import java.util.Collections;
33+
import java.util.HashMap;
3134
import java.util.List;
35+
import java.util.Map;
3236
import java.util.Optional;
3337
import org.testng.annotations.BeforeMethod;
3438
import org.testng.annotations.Test;
@@ -358,4 +362,52 @@ public void testListStoresWithDataReplicationPolicyFilterNullPolicy() {
358362

359363
assertEquals(response.getStoreNamesCount(), 0);
360364
}
365+
366+
@Test
367+
public void testGetClusterHealthStoresSuccess() {
368+
GetClusterHealthStoresGrpcRequest request =
369+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName("testCluster").build();
370+
371+
Map<String, String> storeStatusMap = new HashMap<>();
372+
storeStatusMap.put("store1", "ONLINE");
373+
storeStatusMap.put("store2", "DEGRADED");
374+
storeStatusMap.put("store3", "UNAVAILABLE");
375+
when(admin.getAllStoreStatuses("testCluster")).thenReturn(storeStatusMap);
376+
377+
GetClusterHealthStoresGrpcResponse response = storeRequestHandler.getClusterHealthStores(request);
378+
379+
verify(admin, times(1)).getAllStoreStatuses("testCluster");
380+
assertEquals(response.getClusterName(), "testCluster");
381+
assertEquals(response.getStoreStatusMapCount(), 3);
382+
assertEquals(response.getStoreStatusMapMap().get("store1"), "ONLINE");
383+
assertEquals(response.getStoreStatusMapMap().get("store2"), "DEGRADED");
384+
assertEquals(response.getStoreStatusMapMap().get("store3"), "UNAVAILABLE");
385+
}
386+
387+
@Test
388+
public void testGetClusterHealthStoresEmptyMap() {
389+
GetClusterHealthStoresGrpcRequest request =
390+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName("testCluster").build();
391+
392+
when(admin.getAllStoreStatuses("testCluster")).thenReturn(Collections.emptyMap());
393+
394+
GetClusterHealthStoresGrpcResponse response = storeRequestHandler.getClusterHealthStores(request);
395+
396+
verify(admin, times(1)).getAllStoreStatuses("testCluster");
397+
assertEquals(response.getClusterName(), "testCluster");
398+
assertEquals(response.getStoreStatusMapCount(), 0);
399+
}
400+
401+
@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Cluster name is required")
402+
public void testGetClusterHealthStoresMissingClusterName() {
403+
GetClusterHealthStoresGrpcRequest request = GetClusterHealthStoresGrpcRequest.newBuilder().build();
404+
storeRequestHandler.getClusterHealthStores(request);
405+
}
406+
407+
@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Cluster name is required")
408+
public void testGetClusterHealthStoresEmptyClusterName() {
409+
GetClusterHealthStoresGrpcRequest request =
410+
GetClusterHealthStoresGrpcRequest.newBuilder().setClusterName("").build();
411+
storeRequestHandler.getClusterHealthStores(request);
412+
}
361413
}

0 commit comments

Comments
 (0)