Skip to content

Commit 12992e3

Browse files
emma-sgemma-sg
committed
use causal consistency and remove transaction within quota update
1 parent 617c99a commit 12992e3

File tree

1 file changed

+81
-84
lines changed

1 file changed

+81
-84
lines changed

‎backend/btrixcloud/orgs.py‎

Lines changed: 81 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -635,96 +635,93 @@ async def update_quotas(
635635
) -> None:
636636
"""update organization quotas"""
637637

638-
async with await self.dbclient.start_session() as session:
638+
async with await self.dbclient.start_session(
639+
causal_consistency=True
640+
) as session:
639641
try:
640-
async with session.start_transaction():
641-
# Re-fetch the organization within the context of the transaction
642-
# so that the operation as a whole is atomic.
643-
org = await self.get_org_by_id(org.id, session=session)
644-
645-
previous_extra_mins = (
646-
org.quotas.extraExecMinutes
647-
if (org.quotas and org.quotas.extraExecMinutes)
648-
else 0
649-
)
650-
previous_gifted_mins = (
651-
org.quotas.giftedExecMinutes
652-
if (org.quotas and org.quotas.giftedExecMinutes)
653-
else 0
654-
)
655-
656-
if mode == "add":
657-
increment_update: dict[str, Any] = {
658-
"$inc": {},
659-
}
642+
# Re-fetch the organization within the session
643+
# so that the operation as a whole is atomic.
644+
org = await self.get_org_by_id(org.id, session=session)
645+
646+
previous_extra_mins = (
647+
org.quotas.extraExecMinutes
648+
if (org.quotas and org.quotas.extraExecMinutes)
649+
else 0
650+
)
651+
previous_gifted_mins = (
652+
org.quotas.giftedExecMinutes
653+
if (org.quotas and org.quotas.giftedExecMinutes)
654+
else 0
655+
)
660656

661-
for field, value in quotas.model_dump(
662-
exclude_unset=True, exclude_defaults=True, exclude_none=True
663-
).items():
664-
if field == "context" or value is None:
665-
continue
666-
inc = max(value, -org.quotas.model_dump().get(field, 0))
667-
increment_update["$inc"][f"quotas.{field}"] = inc
668-
669-
updated_org = await self.orgs.find_one_and_update(
670-
{"_id": org.id},
671-
increment_update,
672-
projection={"quotas": True},
673-
return_document=ReturnDocument.AFTER,
674-
session=session,
675-
)
676-
quotas = OrgQuotasIn(**updated_org["quotas"])
677-
678-
update: dict[str, dict[str, dict[str, Any] | int]] = {
679-
"$push": {
680-
"quotaUpdates": OrgQuotaUpdate(
681-
modified=dt_now(),
682-
update=OrgQuotas(
683-
**quotas.model_dump(
684-
exclude_unset=True,
685-
exclude_defaults=True,
686-
exclude_none=True,
687-
)
688-
),
689-
context=context,
690-
).model_dump()
691-
},
657+
if mode == "add":
658+
increment_update: dict[str, Any] = {
692659
"$inc": {},
693-
"$set": {},
694660
}
695661

696-
if mode == "set":
697-
increment_update = quotas.model_dump(
698-
exclude_unset=True, exclude_defaults=True, exclude_none=True
699-
)
700-
update["$set"]["quotas"] = increment_update
701-
702-
# Inc org available fields for extra/gifted execution time as needed
703-
if quotas.extraExecMinutes is not None:
704-
extra_secs_diff = (
705-
quotas.extraExecMinutes - previous_extra_mins
706-
) * 60
707-
if org.extraExecSecondsAvailable + extra_secs_diff <= 0:
708-
update["$set"]["extraExecSecondsAvailable"] = 0
709-
else:
710-
update["$inc"][
711-
"extraExecSecondsAvailable"
712-
] = extra_secs_diff
713-
714-
if quotas.giftedExecMinutes is not None:
715-
gifted_secs_diff = (
716-
quotas.giftedExecMinutes - previous_gifted_mins
717-
) * 60
718-
if org.giftedExecSecondsAvailable + gifted_secs_diff <= 0:
719-
update["$set"]["giftedExecSecondsAvailable"] = 0
720-
else:
721-
update["$inc"][
722-
"giftedExecSecondsAvailable"
723-
] = gifted_secs_diff
724-
725-
await self.orgs.find_one_and_update(
726-
{"_id": org.id}, update, session=session
662+
for field, value in quotas.model_dump(
663+
exclude_unset=True, exclude_defaults=True, exclude_none=True
664+
).items():
665+
if field == "context" or value is None:
666+
continue
667+
inc = max(value, -org.quotas.model_dump().get(field, 0))
668+
increment_update["$inc"][f"quotas.{field}"] = inc
669+
670+
updated_org = await self.orgs.find_one_and_update(
671+
{"_id": org.id},
672+
increment_update,
673+
projection={"quotas": True},
674+
return_document=ReturnDocument.AFTER,
675+
session=session,
727676
)
677+
quotas = OrgQuotasIn(**updated_org["quotas"])
678+
679+
update: dict[str, dict[str, dict[str, Any] | int]] = {
680+
"$push": {
681+
"quotaUpdates": OrgQuotaUpdate(
682+
modified=dt_now(),
683+
update=OrgQuotas(
684+
**quotas.model_dump(
685+
exclude_unset=True,
686+
exclude_defaults=True,
687+
exclude_none=True,
688+
)
689+
),
690+
context=context,
691+
).model_dump()
692+
},
693+
"$inc": {},
694+
"$set": {},
695+
}
696+
697+
if mode == "set":
698+
increment_update = quotas.model_dump(
699+
exclude_unset=True, exclude_defaults=True, exclude_none=True
700+
)
701+
update["$set"]["quotas"] = increment_update
702+
703+
# Inc org available fields for extra/gifted execution time as needed
704+
if quotas.extraExecMinutes is not None:
705+
extra_secs_diff = (
706+
quotas.extraExecMinutes - previous_extra_mins
707+
) * 60
708+
if org.extraExecSecondsAvailable + extra_secs_diff <= 0:
709+
update["$set"]["extraExecSecondsAvailable"] = 0
710+
else:
711+
update["$inc"]["extraExecSecondsAvailable"] = extra_secs_diff
712+
713+
if quotas.giftedExecMinutes is not None:
714+
gifted_secs_diff = (
715+
quotas.giftedExecMinutes - previous_gifted_mins
716+
) * 60
717+
if org.giftedExecSecondsAvailable + gifted_secs_diff <= 0:
718+
update["$set"]["giftedExecSecondsAvailable"] = 0
719+
else:
720+
update["$inc"]["giftedExecSecondsAvailable"] = gifted_secs_diff
721+
722+
await self.orgs.find_one_and_update(
723+
{"_id": org.id}, update, session=session
724+
)
728725
except Exception as e:
729726
print(f"Error updating organization quotas: {e}")
730727
raise HTTPException(status_code=500, detail=str(e)) from e

0 commit comments

Comments
 (0)