Astra multi-tenancy - Shard Assignment #25

zarna1parekh · 2025-02-25T22:39:48Z

Summary

This PR implements 3 GRPC API calls in cluster manager, and updating the existing UpdatePartitionAssignment to handle auto assignment of partitions.

Implementation to manage PartitionMetadata related information is also added. The PartitionMetadata tracks each partitions perPartitionCapacity and dedicated status (shared between tenants or now).

CreatePartition --> Creates / Updates a Partition in ASTRA zookeeper partition_map.

curl \-XPOST \-H 'content-type: application/json; charset=utf-8; protocol=gRPC' http://localhost:8083'/slack.proto.astra.ManagerApiService/CreatePartition' \-d '{
"partition_id": "5"
}'
{"partitionId":"5","maxCapacity":"5000000"}%

ListPartition --> List all the existing partition in the partition_map on zookeeper

curl  -XPOST -H 'content-type: application/json; charset=utf-8; protocol=gRPC'  http://localhost:8083/slack.proto.astra.ManagerApiService/ListPartition -d '{}' 

{
  "partitionMetadata": [
    {
      "partitionId": "5",
      "provisionedCapacity": "1000000",
      "maxCapacity": "5000000",
      "dedicatedPartition": true
    },
    {
      "partitionId": "4",
      "provisionedCapacity": "1000000",
      "maxCapacity": "5000000",
      "dedicatedPartition": true
    },
    {
      "partitionId": "3",
      "provisionedCapacity": "4000000",
      "maxCapacity": "5000000"
    },
    {
      "partitionId": "2",
      "provisionedCapacity": "4000000",
      "maxCapacity": "5000000"
    },
    {
      "partitionId": "1",
      "provisionedCapacity": "4000000",
      "maxCapacity": "5000000"
    }
  ]
}

DeleteDatasetMetadata --> remove a datasetMetadata currently taking traffic from ASTRA. Indirectly, manages the partitions capacity also.

curl \-XPOST \-H 'content-type: application/json; charset=utf-8; protocol=gRPC' http://localhost:8083'/slack.proto.astra.ManagerApiService/DeleteDatasetMetadata' \-d '{  
"name": "tenant1"
}'

{"status": "Deleted dataset metadata tenant1 successfully."}%

UpdatePartitionMetadata --> when no partitionIds are assigned it will auto-assign. Also, added an additional flag require_dedicated_partition if the tenant required dedicated partitions.

curl -XPOST -H 'content-type: application/json; charset=utf-8; protocol=gRPC' http://localhost:8083'/slack.proto.astra.ManagerApiService/UpdatePartitionAssignment' -d '{
          "name": "tenant1",
          "throughputBytes": "8000000",
          "partitionIds": [],
          "requireDedicatedPartition": true}'
{"assignedPartitionIds":["1","2"]}%

Testing

Tested in local dev env by running various APIs.
Added testcases for new APIs and related methods.

Requirements

I've read and understood the Contributing Guidelines and have done my best effort to follow them.
I've read and agree to the Code of Conduct.

mansu · 2025-03-07T21:39:48Z

astra/src/main/java/com/slack/astra/metadata/dataset/DatasetPartitionMetadata.java

+      long startTimeEpochMs,
+      long endTimeEpochMs,
+      List<String> partitions,
+      boolean usingDedicatedPartition) {


Why do we need usingDedicatedPartition field here? I think we can drop it.

I was previously the flag when clear operation fails and need to revert the partitionMap config. But thinking more, I think I can do without it.

mansu · 2025-03-11T16:27:42Z

astra/src/main/java/com/slack/astra/metadata/partition/PartitionMetadata.java

+import com.slack.astra.metadata.core.AstraMetadata;
+import java.util.Objects;
+
+/** PartitionMetadata Object to track the utilization and isPartitionShared in zookeeper */


Please expand on the class doc. PartitionMetadata Object is used to track the utilization of a kafka partition in Zookeeper. isPartitionShared field is set when this partition is shared by multiple tenants. If false, this partition is used by a single tenant.

mansu · 2025-03-11T16:28:35Z

astra/src/main/java/com/slack/astra/metadata/partition/PartitionMetadata.java

+public class PartitionMetadata extends AstraMetadata {
+  public final String partitionId;
+  public final long utilization;
+  public final boolean isPartitionShared;


minor nit but isSharedPartition may be more appropriate.

I had some thoughts on expanding and renaming the members of this Object:

String partitionId boolean dedicatedPartition long provisionedCapacity long maxCapacity

let me know if these fields sound better than existing naming convention.

mansu · 2025-03-11T16:31:55Z