|
@@ -18,6 +18,42 @@ grpc::Status readOnlyStatus(const std::string& rpcName, const DatabaseService& s
|
|
|
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
|
|
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
|
|
|
rpcName + " rejected: database is in read-only mode. " + svc.readOnlyReason());
|
|
rpcName + " rejected: database is in read-only mode. " + svc.readOnlyReason());
|
|
|
}
|
|
}
|
|
|
|
|
+
|
|
|
|
|
+// RAII guard for a streaming-RPC concurrency counter (v1.6.2).
|
|
|
|
|
+// Increments on construction (if under limit), decrements on destruction.
|
|
|
|
|
+// Check ok() before using — false means the limit was reached and the
|
|
|
|
|
+// counter has already been rolled back.
|
|
|
|
|
+class StreamGuard {
|
|
|
|
|
+public:
|
|
|
|
|
+ StreamGuard(std::atomic<uint32_t>& counter, uint32_t max, const char* rpcName)
|
|
|
|
|
+ : counter_(counter), max_(max), rpcName_(rpcName) {
|
|
|
|
|
+ uint32_t prev = counter_.fetch_add(1, std::memory_order_acq_rel);
|
|
|
|
|
+ if (prev >= max_) {
|
|
|
|
|
+ counter_.fetch_sub(1, std::memory_order_acq_rel);
|
|
|
|
|
+ acquired_ = false;
|
|
|
|
|
+ spdlog::warn("gRPC: rejecting {} -- concurrency limit reached ({} >= {})",
|
|
|
|
|
+ rpcName_, prev + 1, max_);
|
|
|
|
|
+ } else {
|
|
|
|
|
+ acquired_ = true;
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ ~StreamGuard() {
|
|
|
|
|
+ if (acquired_) {
|
|
|
|
|
+ counter_.fetch_sub(1, std::memory_order_acq_rel);
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ StreamGuard(const StreamGuard&) = delete;
|
|
|
|
|
+ StreamGuard& operator=(const StreamGuard&) = delete;
|
|
|
|
|
+
|
|
|
|
|
+ bool ok() const { return acquired_; }
|
|
|
|
|
+
|
|
|
|
|
+private:
|
|
|
|
|
+ std::atomic<uint32_t>& counter_;
|
|
|
|
|
+ uint32_t max_;
|
|
|
|
|
+ const char* rpcName_;
|
|
|
|
|
+ bool acquired_ = false;
|
|
|
|
|
+};
|
|
|
|
|
+
|
|
|
} // anonymous namespace
|
|
} // anonymous namespace
|
|
|
|
|
|
|
|
// ===== DatabaseGrpcImpl =====
|
|
// ===== DatabaseGrpcImpl =====
|
|
@@ -838,6 +874,13 @@ grpc::Status DatabaseGrpcImpl::UploadFile(
|
|
|
grpc::ServerReader<pb::FileChunk>* reader,
|
|
grpc::ServerReader<pb::FileChunk>* reader,
|
|
|
pb::UploadFileResponse* response
|
|
pb::UploadFileResponse* response
|
|
|
) {
|
|
) {
|
|
|
|
|
+ // v1.6.2 — cap concurrent file streams to bound memory growth.
|
|
|
|
|
+ StreamGuard guard(activeFileStreams_, maxFileStreams_, "UploadFile");
|
|
|
|
|
+ if (!guard.ok()) {
|
|
|
|
|
+ return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
|
|
|
|
|
+ "too many concurrent file streams (max " +
|
|
|
|
|
+ std::to_string(maxFileStreams_) + " reached)");
|
|
|
|
|
+ }
|
|
|
if (service_.isReadOnly()) {
|
|
if (service_.isReadOnly()) {
|
|
|
return readOnlyStatus("UploadFile", service_);
|
|
return readOnlyStatus("UploadFile", service_);
|
|
|
}
|
|
}
|
|
@@ -884,6 +927,13 @@ grpc::Status DatabaseGrpcImpl::DownloadFile(
|
|
|
const pb::DownloadFileRequest* request,
|
|
const pb::DownloadFileRequest* request,
|
|
|
grpc::ServerWriter<pb::FileChunk>* writer
|
|
grpc::ServerWriter<pb::FileChunk>* writer
|
|
|
) {
|
|
) {
|
|
|
|
|
+ // v1.6.2 — cap concurrent file streams to bound memory growth.
|
|
|
|
|
+ StreamGuard guard(activeFileStreams_, maxFileStreams_, "DownloadFile");
|
|
|
|
|
+ if (!guard.ok()) {
|
|
|
|
|
+ return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
|
|
|
|
|
+ "too many concurrent file streams (max " +
|
|
|
|
|
+ std::to_string(maxFileStreams_) + " reached)");
|
|
|
|
|
+ }
|
|
|
try {
|
|
try {
|
|
|
auto info = files_.getFileInfo(request->id());
|
|
auto info = files_.getFileInfo(request->id());
|
|
|
if (!info) {
|
|
if (!info) {
|
|
@@ -999,6 +1049,13 @@ grpc::Status DatabaseGrpcImpl::Subscribe(
|
|
|
const pb::SubscribeRequest* request,
|
|
const pb::SubscribeRequest* request,
|
|
|
grpc::ServerWriter<pb::DatabaseEvent>* writer
|
|
grpc::ServerWriter<pb::DatabaseEvent>* writer
|
|
|
) {
|
|
) {
|
|
|
|
|
+ // v1.6.2 — cap concurrent Subscribe streams to bound memory growth.
|
|
|
|
|
+ StreamGuard guard(activeSubscribeStreams_, maxSubscribeStreams_, "Subscribe");
|
|
|
|
|
+ if (!guard.ok()) {
|
|
|
|
|
+ return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
|
|
|
|
|
+ "too many concurrent Subscribe streams (max " +
|
|
|
|
|
+ std::to_string(maxSubscribeStreams_) + " reached)");
|
|
|
|
|
+ }
|
|
|
std::vector<std::string> collections(request->collections().begin(), request->collections().end());
|
|
std::vector<std::string> collections(request->collections().begin(), request->collections().end());
|
|
|
std::vector<std::string> patterns(request->patterns().begin(), request->patterns().end());
|
|
std::vector<std::string> patterns(request->patterns().begin(), request->patterns().end());
|
|
|
bool includeData = request->include_data();
|
|
bool includeData = request->include_data();
|