SystemQuotaController implements QuotaController<Void> { @Override public EvictionQuota getQuota(Reference reference) { EvictionQuota.Builder quotaBuilder = EvictionQuota.newBuilder().withReference(reference); if (!inTimeWindowPredicate.get()) { return quotaBuilder.withQuota(0).withMessage(OUTSIDE_SYSTEM_TIME_WINDOW.getRejectionReason().get()).build(); } if (systemTokenBucket.getNumberOfTokens() <= 0) { return quotaBuilder.withQuota(0).withMessage(QUOTA_LIMIT_EXCEEDED.getRejectionReason().get()).build(); } return quotaBuilder.withQuota(systemTokenBucket.getNumberOfTokens()).withMessage(quotaMessage).build(); } @Inject  SystemQuotaController(SystemDisruptionBudgetResolver systemDisruptionBudgetResolver,
                                 TitusRuntime titusRuntime); @VisibleForTesting  SystemQuotaController(SystemDisruptionBudgetResolver systemDisruptionBudgetResolver,
                          Duration retryInterval,
                          TitusRuntime titusRuntime); @PreDestroy void shutdown(); @Override ConsumptionResult consume(String taskId); @Override void giveBackConsumedQuota(String taskId); @Override EvictionQuota getQuota(Reference reference);  }
@Test(timeout = 60_000) public void testRetryOnFailure() { budgetEmitter.onNext(SystemDisruptionBudget.newBasicSystemDisruptionBudget(1, 1)); quotaController = newSystemQuotaController(); budgetEmitter.onError(new RuntimeException("Simulated error")); await().until(this::isBudgetEmitterOpen); budgetEmitter.onNext(SystemDisruptionBudget.newBasicSystemDisruptionBudget(5, 5)); await().until(() -> quotaController.getQuota(Reference.system()).getQuota() == 5); }