Skip to content

Commit a69bc05

Browse files
authored
[ENH] Control how far into the future s3heap scans (#5620)
## Description of changes We need to not fetch more than a minute or two into the future to make s3heap-service performant and cost-effective. This PR makes it so that the s3heap will allow limiting how far into the future a call will look for peek/prune. ## Test plan CI ## Migration plan N/A ## Observability plan N/A ## Documentation Changes N/A
1 parent 15895fd commit a69bc05

File tree

8 files changed

+718
-16
lines changed

8 files changed

+718
-16
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

idl/chromadb/proto/heapservice.proto

Lines changed: 107 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,116 @@ syntax = "proto3";
33
package chroma;
44

55
import "chromadb/proto/chroma.proto";
6+
import "google/protobuf/timestamp.proto";
67

8+
// A task that can be scheduled and triggered in the heap.
9+
message Triggerable {
10+
string partitioning_uuid = 1;
11+
string scheduling_uuid = 2;
12+
}
13+
14+
// A scheduled task with its next execution time and unique identifier.
15+
message Schedule {
16+
Triggerable triggerable = 1;
17+
google.protobuf.Timestamp next_scheduled = 2;
18+
string nonce = 3;
19+
}
20+
21+
// A heap item with its scheduled time.
22+
message HeapItem {
23+
Triggerable triggerable = 1;
24+
string nonce = 2;
25+
google.protobuf.Timestamp scheduled_time = 3;
26+
}
27+
28+
// Limits on range-scan-backed operations.
29+
message Limits {
30+
optional uint32 buckets_to_read = 1;
31+
optional uint32 max_items = 2;
32+
optional google.protobuf.Timestamp time_cut_off = 3;
33+
}
34+
35+
// Statistics from a pruning operation.
36+
message PruneStats {
37+
uint32 items_pruned = 1;
38+
uint32 items_retained = 2;
39+
uint32 buckets_deleted = 3;
40+
uint32 buckets_updated = 4;
41+
}
42+
43+
// Filter criteria for querying heap items.
44+
message FilterCriteria {
45+
optional string partitioning_uuid = 1;
46+
optional string scheduling_uuid = 2;
47+
}
48+
49+
// Request to manually add schedules to the heap.
50+
message PushRequest {
51+
repeated Schedule schedules = 1;
52+
}
53+
54+
// Response from pushing schedules.
55+
message PushResponse {
56+
uint32 schedules_added = 1;
57+
}
58+
59+
// Request to query heap items with filters.
60+
message PeekRequest {
61+
Limits limits = 1;
62+
optional FilterCriteria filter = 2;
63+
}
64+
65+
// Response containing heap items.
66+
message PeekResponse {
67+
repeated HeapItem items = 1;
68+
}
69+
70+
// Request to prune completed tasks.
71+
message PruneRequest {
72+
Limits limits = 1;
73+
}
74+
75+
// Response from pruning operation.
76+
message PruneResponse {
77+
PruneStats stats = 1;
78+
}
79+
80+
// Request to prune a specific bucket.
81+
message PruneBucketRequest {
82+
google.protobuf.Timestamp bucket = 1;
83+
}
84+
85+
// Response from bucket pruning operation.
86+
message PruneBucketResponse {
87+
PruneStats stats = 1;
88+
}
89+
90+
// Request to list buckets in the heap.
91+
message ListBucketsRequest {
92+
optional uint32 max_buckets = 1;
93+
}
94+
95+
// Response containing bucket timestamps.
96+
message ListBucketsResponse {
97+
repeated google.protobuf.Timestamp buckets = 1;
98+
}
99+
100+
// Request for heap summary statistics.
7101
message HeapSummaryRequest {}
8-
message HeapSummaryResponse {}
102+
103+
// Response with heap summary statistics.
104+
message HeapSummaryResponse {
105+
uint32 total_items = 1;
106+
optional google.protobuf.Timestamp oldest_bucket = 2;
107+
optional google.protobuf.Timestamp newest_bucket = 3;
108+
uint32 bucket_count = 4;
109+
}
9110

10111
service HeapTenderService {
112+
rpc Push(PushRequest) returns (PushResponse) {}
113+
rpc Peek(PeekRequest) returns (PeekResponse) {}
114+
rpc Prune(PruneRequest) returns (PruneResponse) {}
115+
rpc PruneBucket(PruneBucketRequest) returns (PruneBucketResponse) {}
116+
rpc ListBuckets(ListBucketsRequest) returns (ListBucketsResponse) {}
11117
rpc Summary(HeapSummaryRequest) returns (HeapSummaryResponse) {}
12118
}

rust/s3heap-service/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ async-trait = { workspace = true }
88
chrono = { workspace = true }
99
figment = { workspace = true }
1010
futures = { workspace = true }
11+
prost-types = { workspace = true }
1112
serde = { workspace = true }
1213
serde_json = { workspace = true }
1314
tokio = { workspace = true }

0 commit comments

Comments
 (0)