PrometheusBasedResourceLimitChecks implements ResourceLimitChecks { @Override public Future<Boolean> isConnectionLimitReached(final TenantObject tenant, final SpanContext spanContext) { Objects.requireNonNull(tenant); final Span span = createSpan("verify connection limit", spanContext, tenant); final Map<String, Object> traceItems = new HashMap<>(); final Promise<Boolean> result = Promise.promise(); if (tenant.getResourceLimits() == null) { traceItems.put(Fields.EVENT, "no resource limits configured"); LOG.trace("no resource limits configured for tenant [{}]", tenant.getTenantId()); result.complete(Boolean.FALSE); } else { final long maxConnections = tenant.getResourceLimits().getMaxConnections(); LOG.trace("connection limit for tenant [{}] is [{}]", tenant.getTenantId(), maxConnections); if (maxConnections == TenantConstants.UNLIMITED_CONNECTIONS) { traceItems.put(Fields.EVENT, "no connection limit configured"); result.complete(Boolean.FALSE); } else { final AtomicBoolean cacheHit = new AtomicBoolean(true); connectionCountCache.get(tenant.getTenantId(), (tenantId, executor) -> { final CompletableFuture<LimitedResource<Long>> r = new CompletableFuture<>(); cacheHit.set(false); final String queryParams = String.format( "sum(%s{tenant=\"%s\"})", METRIC_NAME_CONNECTIONS, tenantId); executeQuery(queryParams, span) .onSuccess(currentConnections -> r.complete(new LimitedResource<Long>(maxConnections, currentConnections))) .onFailure(t -> r.completeExceptionally(t)); return r; }) .whenComplete((value, error) -> { TracingHelper.TAG_CACHE_HIT.set(span, cacheHit.get()); if (error != null) { TracingHelper.logError(span, error); result.complete(Boolean.FALSE); } else { traceItems.put(TenantConstants.FIELD_MAX_CONNECTIONS, maxConnections); traceItems.put("current-connections", value.getCurrentValue()); final boolean isExceeded = value.getCurrentValue() >= value.getCurrentLimit(); LOG.trace("connection limit {}exceeded [tenant: {}, current connections: {}, max-connections: {}]", isExceeded ? "" : "not ", tenant.getTenantId(), value.getCurrentValue(), value.getCurrentLimit()); result.complete(isExceeded); } }); } } return result.future().map(b -> { traceItems.put("limit exceeded", b); span.log(traceItems); 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 testExecuteQuerySetsAuthHeader(final VertxTestContext ctx) { final PrometheusBasedResourceLimitChecksConfig config = new PrometheusBasedResourceLimitChecksConfig(); config.setUsername("hono"); config.setPassword("hono-secret"); limitChecksImpl = new PrometheusBasedResourceLimitChecks( webClient, config, connectionCountCache, connectionDurationCache, dataVolumeCache, tracer); givenCurrentConnections(0); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits().setMaxConnections(10)); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request).basicAuthentication(eq("hono"), eq("hono-secret")); verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testConnectionLimitCheckSucceedsIfNoResourceLimitsAreSet(final VertxTestContext ctx) { givenCurrentConnections(9); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request, never()).send(any(Handler.class)); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testConnectionLimitCheckSucceedsIfUnlimited(final VertxTestContext ctx) { givenCurrentConnections(9); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits().setMaxConnections(TenantConstants.UNLIMITED_CONNECTIONS)); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(request, never()).send(any(Handler.class)); }); ctx.completeNow(); })); }
@SuppressWarnings("unchecked") @Test public void testConnectionLimitCheckSucceeds(final VertxTestContext ctx) { givenCurrentConnections(9); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits().setMaxConnections(10)); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(req).addQueryParam( "query", getExpectedConnectionNumberQuery(tenant)); verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }
@Test @SuppressWarnings("unchecked") public void testConnectionLimitCheckFails(final VertxTestContext ctx) { givenCurrentConnections(10); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits().setMaxConnections(10)); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertTrue(response); verify(req).addQueryParam( "query", getExpectedConnectionNumberQuery(tenant)); verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }
@Test @SuppressWarnings("unchecked") public void testConnectionLimitCheckSucceedsIfQueryTimesOut(final VertxTestContext ctx) { givenFailResponseWithTimeoutException(); final TenantObject tenant = TenantObject.from(Constants.DEFAULT_TENANT, true) .setResourceLimits(new ResourceLimits().setMaxConnections(10)); limitChecksImpl.isConnectionLimitReached(tenant, mock(SpanContext.class)).onComplete( ctx.succeeding(response -> { ctx.verify(() -> { assertFalse(response); verify(req).addQueryParam( "query", getExpectedConnectionNumberQuery(tenant)); verify(request).send(any(Handler.class)); }); ctx.completeNow(); })); }