feat(kv): set_if compare-and-swap for KV (per-app + shared + SDK) (Track A M5)
No atomic CAS existed, so concurrent writers to the same KV key raced. Add set_if(key, expected, new): writes only when the current value equals expected, or (expected absent) only when the key is missing — the primitive for lock-free counters / optimistic concurrency. - KvService + GroupKvService traits gain set_if with a defaulted "unsupported" impl (Noop + other impls untouched); the real services override it. - kv_repo/group_kv_repo: atomic set_if — a single UPDATE ... WHERE value = $expected (JSONB semantic equality) or INSERT ... ON CONFLICT DO NOTHING. Atomic without a transaction; returns whether the write happened. - Services enforce the same value-size cap + (group) writes-authed + row/byte quotas as set; the event fires only on a successful swap. - Executor Rhai SDK: kv::collection(c).set_if(key, expected, new) and the shared_collection variant; Rhai () for expected means "only if absent". Pinned by kv_service/group_kv_service set_if unit tests (CAS + fail-closed) + sdk_kv::kv_set_if_compare_and_swap (through a real script). No migration. docs/sdk-shape.md documents the primitive. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -39,6 +39,18 @@ pub trait GroupKvRepo: Send + Sync {
|
||||
value: serde_json::Value,
|
||||
) -> Result<Option<serde_json::Value>, GroupKvRepoError>;
|
||||
|
||||
/// Compare-and-swap. Atomically writes `new` iff the current value equals
|
||||
/// `expected` (`expected = None` → iff the key is ABSENT). Returns `true`
|
||||
/// when the write happened. One SQL statement — atomic without a tx.
|
||||
async fn set_if(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, GroupKvRepoError>;
|
||||
|
||||
/// Returns the deleted value if present, `None` if the row didn't exist.
|
||||
async fn delete(
|
||||
&self,
|
||||
@@ -143,6 +155,43 @@ impl GroupKvRepo for PostgresGroupKvRepo {
|
||||
Ok(row.and_then(|(v,)| v))
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, GroupKvRepoError> {
|
||||
let affected = match expected {
|
||||
Some(exp) => sqlx::query(
|
||||
"UPDATE group_kv_entries SET value = $4, updated_at = NOW() \
|
||||
WHERE group_id = $1 AND collection = $2 AND key = $3 AND value = $5",
|
||||
)
|
||||
.bind(group_id.into_inner())
|
||||
.bind(collection)
|
||||
.bind(key)
|
||||
.bind(new)
|
||||
.bind(exp)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected(),
|
||||
None => sqlx::query(
|
||||
"INSERT INTO group_kv_entries (group_id, collection, key, value) \
|
||||
VALUES ($1, $2, $3, $4) \
|
||||
ON CONFLICT (group_id, collection, key) DO NOTHING",
|
||||
)
|
||||
.bind(group_id.into_inner())
|
||||
.bind(collection)
|
||||
.bind(key)
|
||||
.bind(new)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected(),
|
||||
};
|
||||
Ok(affected == 1)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
|
||||
@@ -263,6 +263,66 @@ impl GroupKvService for GroupKvServiceImpl {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
cx: &SdkCallCx,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, GroupKvError> {
|
||||
let group_id = self.owning_group(cx, collection).await?;
|
||||
let encoded_len = serde_json::to_vec(&new)
|
||||
.map(|v| v.len())
|
||||
.map_err(|e| GroupKvError::Backend(format!("encode value: {e}")))?;
|
||||
if encoded_len > self.max_value_bytes {
|
||||
return Err(GroupKvError::ValueTooLarge {
|
||||
limit: self.max_value_bytes,
|
||||
actual: encoded_len,
|
||||
});
|
||||
}
|
||||
self.check_write(cx, group_id).await?;
|
||||
// Same quota rules as `set` (best-effort pre-check; the swap itself is
|
||||
// atomic in the repo). A swap that inserts a NEW key (expected absent)
|
||||
// must fit the row ceiling; the byte ceiling uses the projected total.
|
||||
let existing = self.repo.get(group_id, collection, key).await?;
|
||||
if expected.is_none() && existing.is_none() {
|
||||
let count = self.repo.count_rows(group_id).await?;
|
||||
if count >= self.max_rows {
|
||||
return Err(GroupKvError::QuotaExceeded {
|
||||
limit: usize::try_from(self.max_rows).unwrap_or(usize::MAX),
|
||||
actual: usize::try_from(count).unwrap_or(usize::MAX),
|
||||
});
|
||||
}
|
||||
}
|
||||
let old_len = existing
|
||||
.as_ref()
|
||||
.and_then(|v| serde_json::to_vec(v).ok())
|
||||
.map_or(0u64, |b| b.len() as u64);
|
||||
let used = self.repo.total_bytes(group_id).await?;
|
||||
let projected = used.saturating_sub(old_len) + encoded_len as u64;
|
||||
if projected > self.max_total_bytes {
|
||||
return Err(GroupKvError::TotalBytesQuotaExceeded {
|
||||
limit: self.max_total_bytes,
|
||||
actual: projected,
|
||||
});
|
||||
}
|
||||
let op = if expected.is_some() {
|
||||
"update"
|
||||
} else {
|
||||
"insert"
|
||||
};
|
||||
let swapped = self
|
||||
.repo
|
||||
.set_if(group_id, collection, key, expected, new.clone())
|
||||
.await?;
|
||||
if swapped {
|
||||
self.emit_shared(cx, group_id, op, collection, key, Some(new))
|
||||
.await;
|
||||
}
|
||||
Ok(swapped)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
cx: &SdkCallCx,
|
||||
@@ -348,6 +408,27 @@ mod tests {
|
||||
.insert((group_id, collection.to_string(), key.to_string()), value))
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, GroupKvRepoError> {
|
||||
let mut data = self.data.lock().await;
|
||||
let k = (group_id, collection.to_string(), key.to_string());
|
||||
let matches = match (&expected, data.get(&k)) {
|
||||
(None, None) => true,
|
||||
(Some(exp), Some(cur)) => exp == cur,
|
||||
_ => false,
|
||||
};
|
||||
if matches {
|
||||
data.insert(k, new);
|
||||
}
|
||||
Ok(matches)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
@@ -643,4 +724,56 @@ mod tests {
|
||||
let err = kv.get(&cx, "", "k").await.unwrap_err();
|
||||
assert!(matches!(err, GroupKvError::InvalidCollection));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn set_if_is_compare_and_swap_and_writes_fail_closed() {
|
||||
let app = AppId::new();
|
||||
let group = GroupId::new();
|
||||
let mut resolver = FakeResolver::default();
|
||||
resolver.map.insert((app, "catalog".into()), group);
|
||||
let kv = svc(resolver);
|
||||
let cx = cx_with(app, Some(owner()));
|
||||
|
||||
// Insert-if-absent, then swap-if-equal (as the per-app path).
|
||||
assert!(kv
|
||||
.set_if(&cx, "catalog", "k", None, serde_json::json!(1))
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(!kv
|
||||
.set_if(&cx, "catalog", "k", None, serde_json::json!(2))
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(!kv
|
||||
.set_if(
|
||||
&cx,
|
||||
"catalog",
|
||||
"k",
|
||||
Some(serde_json::json!(9)),
|
||||
serde_json::json!(3)
|
||||
)
|
||||
.await
|
||||
.unwrap());
|
||||
assert!(kv
|
||||
.set_if(
|
||||
&cx,
|
||||
"catalog",
|
||||
"k",
|
||||
Some(serde_json::json!(1)),
|
||||
serde_json::json!(3)
|
||||
)
|
||||
.await
|
||||
.unwrap());
|
||||
assert_eq!(
|
||||
kv.get(&cx, "catalog", "k").await.unwrap(),
|
||||
Some(serde_json::json!(3))
|
||||
);
|
||||
|
||||
// A shared-collection CAS still fails closed for an anonymous writer.
|
||||
let anon = cx_with(app, None);
|
||||
let err = kv
|
||||
.set_if(&anon, "catalog", "k", None, serde_json::json!(4))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, GroupKvError::Forbidden));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,6 +39,19 @@ pub trait KvRepo: Send + Sync {
|
||||
value: serde_json::Value,
|
||||
) -> Result<Option<serde_json::Value>, KvRepoError>;
|
||||
|
||||
/// Compare-and-swap. Atomically writes `new` iff the current value equals
|
||||
/// `expected` (`expected = None` → iff the key is ABSENT). Returns `true`
|
||||
/// when the write happened. A single SQL statement, so it is atomic without a
|
||||
/// transaction.
|
||||
async fn set_if(
|
||||
&self,
|
||||
app_id: AppId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, KvRepoError>;
|
||||
|
||||
/// Returns the deleted value if present, `None` if the row didn't
|
||||
/// exist. The caller turns the `bool was-present` part into the
|
||||
/// SDK's return value; the `Option<value>` part feeds the
|
||||
@@ -131,6 +144,45 @@ impl KvRepo for PostgresKvRepo {
|
||||
Ok(row.and_then(|(v,)| v))
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
app_id: AppId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, KvRepoError> {
|
||||
let affected = match expected {
|
||||
// Swap iff the current value matches (JSONB `=` is semantic equality).
|
||||
Some(exp) => sqlx::query(
|
||||
"UPDATE kv_entries SET value = $4, updated_at = NOW() \
|
||||
WHERE app_id = $1 AND collection = $2 AND key = $3 AND value = $5",
|
||||
)
|
||||
.bind(app_id.into_inner())
|
||||
.bind(collection)
|
||||
.bind(key)
|
||||
.bind(new)
|
||||
.bind(exp)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected(),
|
||||
// Insert iff absent.
|
||||
None => sqlx::query(
|
||||
"INSERT INTO kv_entries (app_id, collection, key, value) \
|
||||
VALUES ($1, $2, $3, $4) \
|
||||
ON CONFLICT (app_id, collection, key) DO NOTHING",
|
||||
)
|
||||
.bind(app_id.into_inner())
|
||||
.bind(collection)
|
||||
.bind(key)
|
||||
.bind(new)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
.rows_affected(),
|
||||
};
|
||||
Ok(affected == 1)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
app_id: AppId,
|
||||
|
||||
@@ -186,6 +186,59 @@ impl KvService for KvServiceImpl {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
cx: &SdkCallCx,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, KvError> {
|
||||
validate_collection(collection)?;
|
||||
let encoded_len = serde_json::to_vec(&new)
|
||||
.map(|v| v.len())
|
||||
.map_err(|e| KvError::Backend(format!("encode value: {e}")))?;
|
||||
if encoded_len > self.max_value_bytes {
|
||||
return Err(KvError::ValueTooLarge {
|
||||
limit: self.max_value_bytes,
|
||||
actual: encoded_len,
|
||||
});
|
||||
}
|
||||
self.check_write(cx).await?;
|
||||
// `insert` when the precondition was "absent", else `update`. The op is
|
||||
// only emitted on a successful swap.
|
||||
let op = if expected.is_some() {
|
||||
"update"
|
||||
} else {
|
||||
"insert"
|
||||
};
|
||||
let swapped = self
|
||||
.repo
|
||||
.set_if(cx.app_id, collection, key, expected, new.clone())
|
||||
.await?;
|
||||
if swapped {
|
||||
// Non-transactional with the write (same caveat as `set`).
|
||||
if let Err(e) = self
|
||||
.events
|
||||
.emit(
|
||||
cx,
|
||||
ServiceEvent {
|
||||
source: "kv",
|
||||
op,
|
||||
collection: Some(collection.to_string()),
|
||||
key: Some(key.to_string()),
|
||||
payload: Some(new),
|
||||
old_payload: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(error = %e, source = "kv", op, event_emit_failure = true, "event emit failed");
|
||||
}
|
||||
}
|
||||
Ok(swapped)
|
||||
}
|
||||
|
||||
async fn delete(&self, cx: &SdkCallCx, collection: &str, key: &str) -> Result<bool, KvError> {
|
||||
validate_collection(collection)?;
|
||||
self.check_write(cx).await?;
|
||||
@@ -283,6 +336,28 @@ mod tests {
|
||||
.insert((app_id, collection.to_string(), key.to_string()), value))
|
||||
}
|
||||
|
||||
async fn set_if(
|
||||
&self,
|
||||
app_id: AppId,
|
||||
collection: &str,
|
||||
key: &str,
|
||||
expected: Option<serde_json::Value>,
|
||||
new: serde_json::Value,
|
||||
) -> Result<bool, KvRepoError> {
|
||||
// Holding the map lock makes the compare-and-set atomic in-memory.
|
||||
let mut data = self.data.lock().await;
|
||||
let k = (app_id, collection.to_string(), key.to_string());
|
||||
let matches = match (&expected, data.get(&k)) {
|
||||
(None, None) => true, // insert-if-absent
|
||||
(Some(exp), Some(cur)) => exp == cur, // swap-if-equal
|
||||
_ => false,
|
||||
};
|
||||
if matches {
|
||||
data.insert(k, new);
|
||||
}
|
||||
Ok(matches)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
app_id: AppId,
|
||||
@@ -464,6 +539,63 @@ mod tests {
|
||||
assert!(matches!(err, KvError::InvalidCollection));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn set_if_is_compare_and_swap() {
|
||||
let kv = svc();
|
||||
let cx = anon_cx(AppId::new());
|
||||
// Insert-if-absent: expected () swaps only when the key is missing.
|
||||
assert!(
|
||||
kv.set_if(&cx, "cas", "k", None, serde_json::json!(1))
|
||||
.await
|
||||
.unwrap(),
|
||||
"insert-if-absent must succeed on a missing key"
|
||||
);
|
||||
assert!(
|
||||
!kv.set_if(&cx, "cas", "k", None, serde_json::json!(2))
|
||||
.await
|
||||
.unwrap(),
|
||||
"insert-if-absent must fail once the key exists"
|
||||
);
|
||||
assert_eq!(
|
||||
kv.get(&cx, "cas", "k").await.unwrap(),
|
||||
Some(serde_json::json!(1))
|
||||
);
|
||||
|
||||
// Swap-if-equal: wrong expected → no-op; correct expected → swaps.
|
||||
assert!(
|
||||
!kv.set_if(
|
||||
&cx,
|
||||
"cas",
|
||||
"k",
|
||||
Some(serde_json::json!(99)),
|
||||
serde_json::json!(3)
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
"a mismatched expected must not swap"
|
||||
);
|
||||
assert_eq!(
|
||||
kv.get(&cx, "cas", "k").await.unwrap(),
|
||||
Some(serde_json::json!(1))
|
||||
);
|
||||
assert!(
|
||||
kv.set_if(
|
||||
&cx,
|
||||
"cas",
|
||||
"k",
|
||||
Some(serde_json::json!(1)),
|
||||
serde_json::json!(3)
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
"a matching expected must swap"
|
||||
);
|
||||
assert_eq!(
|
||||
kv.get(&cx, "cas", "k").await.unwrap(),
|
||||
Some(serde_json::json!(3))
|
||||
);
|
||||
}
|
||||
|
||||
/// Load-bearing: a script with `cx.app_id = A` must NOT see
|
||||
/// entries inserted under `cx.app_id = B`. This is the cross-app
|
||||
/// isolation boundary; getting this wrong is a security
|
||||
|
||||
Reference in New Issue
Block a user