fix(personal): make quota reservation atomic

This commit is contained in:
2026-07-12 02:17:55 +08:00
parent de6d24b4fe
commit c80c807283
2 changed files with 87 additions and 18 deletions
@@ -40,32 +40,39 @@ public class PersonalSpaceService {
@Transactional
public long reserve(PersonalOwner owner, long bytes) {
if (bytes < 0) {
if (bytes < 0 || properties.getMaxItems() <= 0 || properties.getMaxSpaceMb() <= 0) {
throw new ServiceException(QUOTA_EXCEEDED);
}
Map<String, Object> space = ensureAndLockSpace(owner);
long spaceId = ((Number) space.get("id")).longValue();
long used = ((Number) space.get("used_bytes")).longValue();
long quota = ((Number) space.get("quota_bytes")).longValue();
int count = ((Number) space.get("item_count")).intValue();
if (used < 0 || quota < 0 || used > quota || bytes > quota - used
if (used < 0 || quota <= 0 || count < 0 || used > quota || bytes > quota - used
|| count >= properties.getMaxItems()) {
throw new ServiceException(QUOTA_EXCEEDED);
}
return ((Number) space.get("id")).longValue();
int updated = jdbcTemplate.update("""
update aihr_personal_space
set used_bytes = used_bytes + ?, item_count = item_count + 1
where id = ? and tenant_id = ? and owner_user_id = ?
""", bytes, spaceId, owner.tenantId(), owner.userId());
if (updated != 1) {
throw new ServiceException(QUOTA_EXCEEDED);
}
return spaceId;
}
private Map<String, Object> ensureAndLockSpace(PersonalOwner owner) {
try {
return lockSpace(owner);
} catch (EmptyResultDataAccessException ex) {
jdbcTemplate.update("""
insert ignore into aihr_personal_space
(tenant_id, owner_user_id, owner_ext_party_id, quota_bytes)
values (?, ?, ?, ?)
""", owner.tenantId(), owner.userId(), owner.extPartyId(), defaultQuotaBytes());
return lockSpace(owner);
}
jdbcTemplate.update("""
insert into aihr_personal_space
(tenant_id, owner_user_id, owner_ext_party_id, quota_bytes)
values (?, ?, ?, ?)
on duplicate key update id = id
""", owner.tenantId(), owner.userId(), owner.extPartyId(), defaultQuotaBytes());
return lockSpace(owner);
}
private Map<String, Object> lockSpace(PersonalOwner owner) {
@@ -8,6 +8,7 @@ import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.springframework.dao.EmptyResultDataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.mockito.InOrder;
import java.util.Map;
@@ -17,6 +18,7 @@ import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.contains;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -74,19 +76,39 @@ class PersonalSpaceServiceTest {
}
@Test
void reserveCreatesMissingOwnerSpaceWithConfiguredQuotaThenLocksIt() {
void reserveUpsertsOwnerSpaceBeforeLockAndAtomicallyConsumesCapacity() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
when(jdbc.queryForMap(contains("from aihr_personal_space"), eq("000000"), eq(101L)))
.thenThrow(new EmptyResultDataAccessException(1))
.thenReturn(space(7L, 500L * 1024 * 1024, 0L, 0));
when(jdbc.update(contains("set used_bytes = used_bytes + ?"),
eq(1024L), eq(7L), eq("000000"), eq(101L))).thenReturn(1);
PersonalSpaceService service = new PersonalSpaceService(jdbc, properties());
assertEquals(7L, service.reserve(new PersonalOwner("000000", 101L, "ext-101"), 1024L));
verify(jdbc).update(contains("insert ignore into aihr_personal_space"),
InOrder order = inOrder(jdbc);
order.verify(jdbc).update(contains("on duplicate key update"),
eq("000000"), eq(101L), eq("ext-101"), eq(500L * 1024 * 1024));
verify(jdbc, org.mockito.Mockito.times(2)).queryForMap(
contains("for update"), eq("000000"), eq(101L));
order.verify(jdbc).queryForMap(contains("for update"), eq("000000"), eq(101L));
order.verify(jdbc).update(
contains("where id = ? and tenant_id = ? and owner_user_id = ?"),
eq(1024L), eq(7L), eq("000000"), eq(101L));
verifyNoMoreInteractions(jdbc);
}
@Test
void reserveZeroBytesStillConsumesOneItemSlot() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
when(jdbc.queryForMap(contains("from aihr_personal_space"), eq("000000"), eq(101L)))
.thenReturn(space(7L, 1024L, 500L, 3));
when(jdbc.update(contains("set used_bytes = used_bytes + ?"),
eq(0L), eq(7L), eq("000000"), eq(101L))).thenReturn(1);
PersonalSpaceService service = new PersonalSpaceService(jdbc, properties());
assertEquals(7L, service.reserve(new PersonalOwner("000000", 101L, null), 0L));
verify(jdbc).update(contains("item_count = item_count + 1"),
eq(0L), eq(7L), eq("000000"), eq(101L));
}
@Test
@@ -112,6 +134,8 @@ class PersonalSpaceServiceTest {
() -> service.reserve(new PersonalOwner("000000", 101L, null), 2L));
assertEquals("PERSONAL_SPACE_QUOTA_EXCEEDED", error.getMessage());
verify(jdbc).update(contains("on duplicate key update"),
eq("000000"), eq(101L), eq(null), eq(500L * 1024 * 1024));
verify(jdbc).queryForMap(contains("for update"), eq("000000"), eq(101L));
verifyNoMoreInteractions(jdbc);
}
@@ -127,10 +151,48 @@ class PersonalSpaceServiceTest {
() -> service.reserve(new PersonalOwner("000000", 101L, null), 1L));
assertEquals("PERSONAL_SPACE_QUOTA_EXCEEDED", error.getMessage());
verify(jdbc).update(contains("on duplicate key update"),
eq("000000"), eq(101L), eq(null), eq(500L * 1024 * 1024));
verify(jdbc).queryForMap(contains("for update"), eq("000000"), eq(101L));
verifyNoMoreInteractions(jdbc);
}
@Test
void reserveRejectsInvalidSpaceStateWithoutUpdatingCounters() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
when(jdbc.queryForMap(contains("from aihr_personal_space"), eq("000000"), eq(101L)))
.thenReturn(space(7L, 1024L, 0L, -1));
PersonalSpaceService service = new PersonalSpaceService(jdbc, properties());
ServiceException error = assertThrows(ServiceException.class,
() -> service.reserve(new PersonalOwner("000000", 101L, null), 1L));
assertEquals("PERSONAL_SPACE_QUOTA_EXCEEDED", error.getMessage());
verify(jdbc).update(contains("on duplicate key update"),
eq("000000"), eq(101L), eq(null), eq(500L * 1024 * 1024));
verify(jdbc).queryForMap(contains("for update"), eq("000000"), eq(101L));
verifyNoMoreInteractions(jdbc);
}
@Test
void reserveRejectsInvalidConfigurationBeforeTouchingStorage() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalKnowledgeProperties properties = properties();
properties.setMaxItems(0);
PersonalSpaceService service = new PersonalSpaceService(jdbc, properties);
ServiceException invalidItems = assertThrows(ServiceException.class,
() -> service.reserve(new PersonalOwner("000000", 101L, null), 1L));
assertEquals("PERSONAL_SPACE_QUOTA_EXCEEDED", invalidItems.getMessage());
properties.setMaxItems(1000);
properties.setMaxSpaceMb(0);
ServiceException invalidSpace = assertThrows(ServiceException.class,
() -> service.reserve(new PersonalOwner("000000", 101L, null), 1L));
assertEquals("PERSONAL_SPACE_QUOTA_EXCEEDED", invalidSpace.getMessage());
verifyNoInteractions(jdbc);
}
private static Map<String, Object> space(long id, long quota, long used, int count) {
return Map.of(
"id", id,