api: throttle streaming admin read polls
This commit is contained in:
a.tolmachev
+129
-2
@@ -231,10 +231,11 @@ mod tests {
|
||||
AsyncJobHandle, AsyncJobId, DescriptorId, ExecutionConfig, ExecutionMode,
|
||||
GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod, JobStatus, MembershipRole,
|
||||
OperationId, Protocol, RestTarget, SecretKind, SoapBindingStyle, SoapOperationMetadata,
|
||||
SoapTarget, SoapVersion, Target, ToolDescription, TransportBehavior, WorkspaceId,
|
||||
SoapTarget, SoapVersion, StreamSession, StreamStatus, Target, ToolDescription,
|
||||
TransportBehavior, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::{MappingRule, MappingSet};
|
||||
use crank_registry::{CreateAsyncJobRequest, PostgresRegistry};
|
||||
use crank_registry::{CreateAsyncJobRequest, CreateStreamSessionRequest, PostgresRegistry};
|
||||
use crank_runtime::SecretCrypto;
|
||||
use crank_schema::{Schema, SchemaKind};
|
||||
use serde_json::{Value, json};
|
||||
@@ -1810,6 +1811,132 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn rate_limits_rapid_stream_session_reads() {
|
||||
let registry = test_registry().await;
|
||||
let storage_root = test_storage_root("stream_session_rate_limit");
|
||||
let server_addr = grpc_test_support::spawn_unary_echo_server().await;
|
||||
let base_url = spawn_admin_api(build_test_app(registry.clone(), storage_root)).await;
|
||||
let client = authorized_client(&base_url).await;
|
||||
|
||||
let created = client
|
||||
.post(format!("{base_url}/operations"))
|
||||
.json(&test_grpc_session_operation_payload(
|
||||
&server_addr,
|
||||
"echo_session_rate_limit",
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap();
|
||||
let operation_id = OperationId::new(created["operation_id"].as_str().unwrap().to_owned());
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let session_id = crank_core::StreamSessionId::new("sess_rate_limited".to_owned());
|
||||
registry
|
||||
.create_stream_session(CreateStreamSessionRequest {
|
||||
session: &StreamSession {
|
||||
id: session_id.clone(),
|
||||
workspace_id: WorkspaceId::new(DEFAULT_WORKSPACE_ID),
|
||||
agent_id: None,
|
||||
operation_id,
|
||||
protocol: Protocol::Grpc,
|
||||
mode: ExecutionMode::Session,
|
||||
status: StreamStatus::Running,
|
||||
cursor: None,
|
||||
state: json!({ "items": [] }),
|
||||
expires_at: now + time::Duration::minutes(5),
|
||||
last_poll_at: Some(now),
|
||||
created_at: now,
|
||||
closed_at: None,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = client
|
||||
.get(format!("{base_url}/stream-sessions/{}", session_id.as_str()))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let body = response.json::<Value>().await.unwrap();
|
||||
|
||||
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(body["error"]["code"], "rate_limited");
|
||||
assert_eq!(
|
||||
body["error"]["message"],
|
||||
"stream session poll rate limit exceeded"
|
||||
);
|
||||
assert_eq!(body["error"]["context"]["session_id"], session_id.as_str());
|
||||
assert!(body["error"]["context"]["poll_after_ms"].as_u64().unwrap() > 0);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn rate_limits_rapid_async_job_result_polls() {
|
||||
let registry = test_registry().await;
|
||||
let storage_root = test_storage_root("async_job_rate_limit");
|
||||
let upstream_base_url = spawn_upstream_server().await;
|
||||
let base_url = spawn_admin_api(build_test_app(registry.clone(), storage_root)).await;
|
||||
let client = authorized_client(&base_url).await;
|
||||
|
||||
let created = client
|
||||
.post(format!("{base_url}/operations"))
|
||||
.json(&test_rest_async_job_operation_payload(
|
||||
&upstream_base_url,
|
||||
"crm_async_job_rate_limit",
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json::<Value>()
|
||||
.await
|
||||
.unwrap();
|
||||
let operation_id = OperationId::new(created["operation_id"].as_str().unwrap().to_owned());
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let job_id = AsyncJobId::new("job_rate_limited".to_owned());
|
||||
registry
|
||||
.create_async_job(CreateAsyncJobRequest {
|
||||
job: &AsyncJobHandle {
|
||||
id: job_id.clone(),
|
||||
workspace_id: WorkspaceId::new(DEFAULT_WORKSPACE_ID),
|
||||
agent_id: None,
|
||||
operation_id,
|
||||
status: JobStatus::Running,
|
||||
progress: json!({ "pct": 25 }),
|
||||
result: None,
|
||||
error: None,
|
||||
expires_at: Some(now + time::Duration::minutes(5)),
|
||||
last_poll_at: Some(now),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
finished_at: None,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = client
|
||||
.get(format!("{base_url}/async-jobs/{}/result", job_id.as_str()))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let body = response.json::<Value>().await.unwrap();
|
||||
|
||||
assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(body["error"]["code"], "rate_limited");
|
||||
assert_eq!(
|
||||
body["error"]["message"],
|
||||
"async job poll rate limit exceeded"
|
||||
);
|
||||
assert_eq!(body["error"]["context"]["job_id"], job_id.as_str());
|
||||
assert!(body["error"]["context"]["poll_after_ms"].as_u64().unwrap() > 0);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn returns_structured_context_for_missing_stream_session() {
|
||||
|
||||
@@ -1585,6 +1585,9 @@ impl AdminService {
|
||||
)
|
||||
})?;
|
||||
ensure_stream_session_workspace(&session, workspace_id)?;
|
||||
let session = self
|
||||
.touch_stream_session_poll_if_ready(workspace_id, session)
|
||||
.await?;
|
||||
|
||||
Ok(stream_session_detail_view(session))
|
||||
}
|
||||
@@ -1665,6 +1668,11 @@ impl AdminService {
|
||||
)
|
||||
})?;
|
||||
ensure_async_job_workspace(&job, workspace_id)?;
|
||||
let job = if matches!(job.status, JobStatus::Completed) {
|
||||
job
|
||||
} else {
|
||||
self.touch_async_job_poll_if_ready(workspace_id, job).await?
|
||||
};
|
||||
|
||||
Ok(async_job_detail_view(job))
|
||||
}
|
||||
@@ -1710,6 +1718,7 @@ impl AdminService {
|
||||
)
|
||||
})?;
|
||||
ensure_async_job_workspace(&job, workspace_id)?;
|
||||
let job = self.touch_async_job_poll_if_ready(workspace_id, job).await?;
|
||||
|
||||
match job.status {
|
||||
JobStatus::Completed => Ok(job.result.unwrap_or(Value::Null)),
|
||||
@@ -1745,6 +1754,102 @@ impl AdminService {
|
||||
}
|
||||
}
|
||||
|
||||
async fn touch_stream_session_poll_if_ready(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
session: StreamSession,
|
||||
) -> Result<StreamSession, ApiError> {
|
||||
let poll_interval_ms = self
|
||||
.streaming_poll_interval_ms(workspace_id, &session.operation_id)
|
||||
.await?;
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let remaining_delay_ms = session.remaining_poll_delay_ms(now, poll_interval_ms);
|
||||
|
||||
if remaining_delay_ms > 0 {
|
||||
return Err(ApiError::rate_limited_with_context(
|
||||
"stream session poll rate limit exceeded",
|
||||
json!({
|
||||
"session_id": session.id.as_str(),
|
||||
"poll_after_ms": remaining_delay_ms,
|
||||
}),
|
||||
));
|
||||
}
|
||||
|
||||
self.registry
|
||||
.touch_stream_session_poll(&session.id, &now)
|
||||
.await
|
||||
.map_err(ApiError::from)
|
||||
}
|
||||
|
||||
async fn touch_async_job_poll_if_ready(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
job: AsyncJobHandle,
|
||||
) -> Result<AsyncJobHandle, ApiError> {
|
||||
let poll_interval_ms = self
|
||||
.streaming_poll_interval_ms(workspace_id, &job.operation_id)
|
||||
.await?;
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let remaining_delay_ms = job.remaining_poll_delay_ms(now, poll_interval_ms);
|
||||
|
||||
if remaining_delay_ms > 0 {
|
||||
return Err(ApiError::rate_limited_with_context(
|
||||
"async job poll rate limit exceeded",
|
||||
json!({
|
||||
"job_id": job.id.as_str(),
|
||||
"poll_after_ms": remaining_delay_ms,
|
||||
}),
|
||||
));
|
||||
}
|
||||
|
||||
self.registry
|
||||
.touch_async_job_poll(&job.id, &now)
|
||||
.await
|
||||
.map_err(ApiError::from)
|
||||
}
|
||||
|
||||
async fn streaming_poll_interval_ms(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
operation_id: &OperationId,
|
||||
) -> Result<u64, ApiError> {
|
||||
let summary = self
|
||||
.registry
|
||||
.get_operation_summary(workspace_id, operation_id)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
ApiError::not_found_with_context(
|
||||
format!("operation {} was not found", operation_id.as_str()),
|
||||
json!({ "operation_id": operation_id.as_str() }),
|
||||
)
|
||||
})?;
|
||||
let version = self
|
||||
.registry
|
||||
.get_operation_version(workspace_id, operation_id, summary.current_draft_version)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
ApiError::not_found_with_context(
|
||||
format!(
|
||||
"operation version {} for {} was not found",
|
||||
summary.current_draft_version,
|
||||
operation_id.as_str()
|
||||
),
|
||||
json!({
|
||||
"operation_id": operation_id.as_str(),
|
||||
"version": summary.current_draft_version,
|
||||
}),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(version
|
||||
.snapshot
|
||||
.execution_config
|
||||
.streaming
|
||||
.as_ref()
|
||||
.and_then(|streaming| streaming.poll_interval_ms)
|
||||
.unwrap_or(1_000))
|
||||
}
|
||||
|
||||
pub async fn list_operations(
|
||||
&self,
|
||||
workspace_id: &WorkspaceId,
|
||||
@@ -2354,7 +2459,7 @@ impl AdminService {
|
||||
batch_size,
|
||||
}),
|
||||
expires_at,
|
||||
last_poll_at: Some(created_at),
|
||||
last_poll_at: None,
|
||||
created_at,
|
||||
closed_at: None,
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user