PrometheusBasedResourceLimitChecks implements ResourceLimitChecks { @Override public Future<Boolean> isMessageLimitReached( final TenantObject tenant, final long payloadSize, final SpanContext spanContext) { Objects.requireNonNull(tenant); final Span span = createSpan("verify message limit", spanContext, tenant); final Map<String, Object> items = new HashMap<>(); items.put("payload-size", payloadSize); final Promise<Boolean> result = Promise.promise(); if (tenant.getResourceLimits() == null) { items.put(Fields.EVENT, "no resource limits configured"); LOG.trace("no resource limits configured for tenant [{}]", tenant.getTenantId()); result.complete(Boolean.FALSE); } else if (tenant.getResourceLimits().getDataVolume() == null) { items.put(Fields.EVENT, "no message limits configured"); LOG.trace("no message limits configured for tenant [{}]", tenant.getTenantId()); result.complete(Boolean.FALSE); } else { checkMessageLimit(tenant, payloadSize, items, span, result); } return result.future() .map(b -> { items.put("limit exceeded", b); span.log(items); span.finish(); return b; }); } PrometheusBasedResourceLimitChecks(
            final WebClient webClient,
            final PrometheusBasedResourceLimitChecksConfig config,
            final AsyncCache<String, LimitedResource<Long>> connectionCountCache,
            final AsyncCache<String, LimitedResource<Duration>> connectionDurationCache,
            final AsyncCache<String, LimitedResource<Long>> dataVolumeCache);  PrometheusBasedResourceLimitChecks(
            final WebClient webClient,
            final PrometheusBasedResourceLimitChecksConfig config,
            final AsyncCache<String, LimitedResource<Long>> connectionCountCache,
            final AsyncCache<String, LimitedResource<Duration>> connectionDurationCache,
            final AsyncCache<String, LimitedResource<Long>> dataVolumeCache,
            final Tracer tracer); @Override Future<Boolean> isConnectionLimitReached(final TenantObject tenant, final SpanContext spanContext); @Override Future<Boolean> isMessageLimitReached(
            final TenantObject tenant,
            final long payloadSize,
            final SpanContext spanContext); @Override Future<Boolean> isConnectionDurationLimitReached(
            final TenantObject tenant,
            final SpanContext spanContext);  }
@SuppressWarnings("unchecked") @Test public void testMessageLimitNotExceededForMissingMetrics(final VertxTestContext ctx) { givenDataVolumeUsageInBytes(null); final long incomingMessageSize = 20; final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits() .setDataVolume(new DataVolume( Instant.parse("2019-01-03T14:30:00Z"), new ResourceLimitsPeriod(ResourceLimitsPeriod.PERIOD_MODE_DAYS).setNoOfDays(30), 100L))); limitChecksImpl.isMessageLimitReached(tenant, incomingMessageSize, spanContext) .onComplete(ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request).send(any(Handler.class)); verify(span).log(argThat((Map<String, ?> map) -> !"error".equals(map.get(Fields.EVENT)))); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testMessageLimitNotExceededWhenNotConfigured(final VertxTestContext ctx) { final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true); limitChecksImpl.isMessageLimitReached(tenant, 10, mock(SpanContext.class)) .onComplete(ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request, never()).send(any(Handler.class)); }); ctx.completeNow(); })); }
@Test @SuppressWarnings("unchecked") public void testMessageLimitUsesValueFromCache(final VertxTestContext ctx) { when(dataVolumeCache.get(anyString(), any(BiFunction.class))) .then(invocation -> { final CompletableFuture<LimitedResource<Long>> result = new CompletableFuture<>(); result.complete(new LimitedResource<Long>(60L, 100L)); return result; }); final long incomingMessageSize = 20; final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits() .setDataVolume(new DataVolume( Instant.parse("2019-01-03T14:30:00Z"), new ResourceLimitsPeriod(ResourceLimitsPeriod.PERIOD_MODE_DAYS).setNoOfDays(30), 100L))); limitChecksImpl.isMessageLimitReached(tenant, incomingMessageSize, spanContext) .onComplete(ctx.succeeding(exceeded -> { ctx.verify(() -> { assertTrue(exceeded); verify(request, never()).send(any(Handler.class)); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testMessageLimitStoresValueToCache(final VertxTestContext ctx) { givenDataVolumeUsageInBytes(100); final long incomingMessageSize = 20; final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits() .setDataVolume(new DataVolume( Instant.parse("2019-01-03T14:30:00Z"), new ResourceLimitsPeriod(ResourceLimitsPeriod.PERIOD_MODE_DAYS).setNoOfDays(30), 100L))); when(dataVolumeCache.get(anyString(), any(BiFunction.class))).then(invocation -> { final BiFunction<String, Executor, CompletableFuture<LimitedResource<Long>>> provider = invocation.getArgument(1); return provider.apply(invocation.getArgument(0), mock(Executor.class)); }); limitChecksImpl.isMessageLimitReached(tenant, incomingMessageSize, mock(SpanContext.class)) .onComplete(ctx.succeeding(response -> { ctx.verify(() -> { verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testMessageLimitNotExceededWhenTimeoutOccurred(final VertxTestContext ctx) { givenFailResponseWithTimeoutException(); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits() .setDataVolume(new DataVolume( Instant.parse("2019-01-03T14:30:00Z"), new ResourceLimitsPeriod(ResourceLimitsPeriod.PERIOD_MODE_MONTHLY), 50L))); limitChecksImpl.isMessageLimitReached(tenant, 100L, mock(SpanContext.class)) .onComplete(ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }