Defer actions taken after rate limit checks

This commit is contained in:
Jon Chambers 2023-08-18 13:22:57 -04:00 committed by Chris Eager
parent 22905fa8ee
commit ff1ef90a6d
4 changed files with 8 additions and 2 deletions

View File

@ -34,7 +34,7 @@ public class CallingGrpcService extends ReactorCallingGrpc.CallingImplBase {
final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice(); final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
return rateLimiters.getTurnLimiter().validateReactive(authenticatedDevice.accountIdentifier()) return rateLimiters.getTurnLimiter().validateReactive(authenticatedDevice.accountIdentifier())
.then(Mono.fromSupplier(() -> turnTokenGenerator.generate(authenticatedDevice.accountIdentifier()))) .then(Mono.defer(() -> Mono.fromSupplier(() -> turnTokenGenerator.generate(authenticatedDevice.accountIdentifier()))))
.map(turnToken -> GetTurnCredentialsResponse.newBuilder() .map(turnToken -> GetTurnCredentialsResponse.newBuilder()
.setUsername(turnToken.username()) .setUsername(turnToken.username())
.setPassword(turnToken.password()) .setPassword(turnToken.password())

View File

@ -135,7 +135,7 @@ public class KeysGrpcService extends ReactorKeysGrpc.KeysImplBase {
request.getDeviceId(); request.getDeviceId();
return rateLimiters.getPreKeysLimiter().validateReactive(rateLimitKey) return rateLimiters.getPreKeysLimiter().validateReactive(rateLimitKey)
.then(Mono.fromFuture(accountsManager.getByServiceIdentifierAsync(targetIdentifier)) .then(Mono.fromFuture(() -> accountsManager.getByServiceIdentifierAsync(targetIdentifier))
.flatMap(Mono::justOrEmpty)) .flatMap(Mono::justOrEmpty))
.switchIfEmpty(Mono.error(Status.NOT_FOUND.asException())) .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
.flatMap(targetAccount -> .flatMap(targetAccount ->

View File

@ -12,6 +12,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never; import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import io.grpc.ServerInterceptors; import io.grpc.ServerInterceptors;
@ -102,5 +103,7 @@ class CallingGrpcServiceTest {
assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode()); assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode());
assertNotNull(exception.getTrailers()); assertNotNull(exception.getTrailers());
assertEquals(retryAfter, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY)); assertEquals(retryAfter, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY));
verifyNoInteractions(turnTokenGenerator);
} }
} }

View File

@ -14,6 +14,7 @@ import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
import com.google.protobuf.ByteString; import com.google.protobuf.ByteString;
@ -669,5 +670,7 @@ class KeysGrpcServiceTest {
assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode()); assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode());
assertNotNull(exception.getTrailers()); assertNotNull(exception.getTrailers());
assertEquals(retryAfterDuration, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY)); assertEquals(retryAfterDuration, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY));
verifyNoInteractions(accountsManager);
} }
} }