Examples with QueryCapacityExceededException org.apache.druid.query.QueryCapacityExceededException used on opensource projects

Example 1 with QueryCapacityExceededException

use of org.apache.druid.query.QueryCapacityExceededException in project druid by druid-io.

the class SqlResource method doPost.

@POST
@Produces(MediaType.APPLICATION_JSON)
@Consumes(MediaType.APPLICATION_JSON)
public Response doPost(final SqlQuery sqlQuery, @Context final HttpServletRequest req) throws IOException {
    final SqlLifecycle lifecycle = sqlLifecycleFactory.factorize();
    final String sqlQueryId = lifecycle.initialize(sqlQuery.getQuery(), sqlQuery.getContext());
    final String remoteAddr = req.getRemoteAddr();
    final String currThreadName = Thread.currentThread().getName();
    try {
        Thread.currentThread().setName(StringUtils.format("sql[%s]", sqlQueryId));
        lifecycle.setParameters(sqlQuery.getParameterList());
        lifecycle.validateAndAuthorize(req);
        // must add after lifecycle is authorized
        sqlLifecycleManager.add(sqlQueryId, lifecycle);
        lifecycle.plan();
        final SqlRowTransformer rowTransformer = lifecycle.createRowTransformer();
        final Sequence<Object[]> sequence = lifecycle.execute();
        final Yielder<Object[]> yielder0 = Yielders.each(sequence);
        try {
            final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
                Exception e = null;
                CountingOutputStream os = new CountingOutputStream(outputStream);
                Yielder<Object[]> yielder = yielder0;
                try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
                    writer.writeResponseStart();
                    if (sqlQuery.includeHeader()) {
                        writer.writeHeader(rowTransformer.getRowType(), sqlQuery.includeTypesHeader(), sqlQuery.includeSqlTypesHeader());
                    }
                    while (!yielder.isDone()) {
                        final Object[] row = yielder.get();
                        writer.writeRowStart();
                        for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
                            final Object value = rowTransformer.transform(row, i);
                            writer.writeRowField(rowTransformer.getFieldList().get(i), value);
                        }
                        writer.writeRowEnd();
                        yielder = yielder.next(null);
                    }
                    writer.writeResponseEnd();
                } catch (Exception ex) {
                    e = ex;
                    log.error(ex, "Unable to send SQL response [%s]", sqlQueryId);
                    throw new RuntimeException(ex);
                } finally {
                    yielder.close();
                    endLifecycle(sqlQueryId, lifecycle, e, remoteAddr, os.getCount());
                }
            }).header(SQL_QUERY_ID_RESPONSE_HEADER, sqlQueryId);
            if (sqlQuery.includeHeader()) {
                responseBuilder.header(SQL_HEADER_RESPONSE_HEADER, SQL_HEADER_VALUE);
            }
            return responseBuilder.build();
        } catch (Throwable e) {
            // make sure to close yielder if anything happened before starting to serialize the response.
            yielder0.close();
            throw new RuntimeException(e);
        }
    } catch (QueryCapacityExceededException cap) {
        endLifecycle(sqlQueryId, lifecycle, cap, remoteAddr, -1);
        return buildNonOkResponse(QueryCapacityExceededException.STATUS_CODE, cap, sqlQueryId);
    } catch (QueryUnsupportedException unsupported) {
        endLifecycle(sqlQueryId, lifecycle, unsupported, remoteAddr, -1);
        return buildNonOkResponse(QueryUnsupportedException.STATUS_CODE, unsupported, sqlQueryId);
    } catch (QueryTimeoutException timeout) {
        endLifecycle(sqlQueryId, lifecycle, timeout, remoteAddr, -1);
        return buildNonOkResponse(QueryTimeoutException.STATUS_CODE, timeout, sqlQueryId);
    } catch (SqlPlanningException | ResourceLimitExceededException e) {
        endLifecycle(sqlQueryId, lifecycle, e, remoteAddr, -1);
        return buildNonOkResponse(BadQueryException.STATUS_CODE, e, sqlQueryId);
    } catch (ForbiddenException e) {
        endLifecycleWithoutEmittingMetrics(sqlQueryId, lifecycle);
        throw (ForbiddenException) serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(// let ForbiddenExceptionMapper handle this
        e);
    } catch (RelOptPlanner.CannotPlanException e) {
        endLifecycle(sqlQueryId, lifecycle, e, remoteAddr, -1);
        SqlPlanningException spe = new SqlPlanningException(SqlPlanningException.PlanningError.UNSUPPORTED_SQL_ERROR, e.getMessage());
        return buildNonOkResponse(BadQueryException.STATUS_CODE, spe, sqlQueryId);
    }// calcite throws a java.lang.AssertionError which is type error not exception. using throwable will catch all
     catch (Throwable e) {
        log.warn(e, "Failed to handle query: %s", sqlQuery);
        endLifecycle(sqlQueryId, lifecycle, e, remoteAddr, -1);
        return buildNonOkResponse(Status.INTERNAL_SERVER_ERROR.getStatusCode(), QueryInterruptedException.wrapIfNeeded(e), sqlQueryId);
    } finally {
        Thread.currentThread().setName(currThreadName);
    }
}

Example 2 with QueryCapacityExceededException

use of org.apache.druid.query.QueryCapacityExceededException in project druid by druid-io.

the class QueryScheduler method acquireLanes.

/**
 * Acquire a semaphore for both the 'total' and a lane, if any is associated with a query
 */
@VisibleForTesting
List<Bulkhead> acquireLanes(Query<?> query) {
    final String lane = QueryContexts.getLane(query);
    final Optional<BulkheadConfig> laneConfig = lane == null ? Optional.empty() : laneRegistry.getConfiguration(lane);
    final Optional<BulkheadConfig> totalConfig = laneRegistry.getConfiguration(TOTAL);
    List<Bulkhead> hallPasses = new ArrayList<>(2);
    try {
        // if we have a lane, get it first
        laneConfig.ifPresent(config -> {
            Bulkhead laneLimiter = laneRegistry.bulkhead(lane, config);
            if (!laneLimiter.tryAcquirePermission()) {
                throw new QueryCapacityExceededException(lane, config.getMaxConcurrentCalls());
            }
            hallPasses.add(laneLimiter);
        });
        // everyone needs to take one from the total lane; to ensure we don't acquire a lane and never release it, we want
        // to check for total capacity exceeded and release the lane (if present) before throwing capacity exceeded
        // note that this isn't strictly fair: the bulkhead doesn't use a fair semaphore, the first to acquire the lane
        // might lose to one that came after it when acquiring the total, or an unlaned query might lose to a laned query
        totalConfig.ifPresent(config -> {
            Bulkhead totalLimiter = laneRegistry.bulkhead(TOTAL, config);
            if (!totalLimiter.tryAcquirePermission()) {
                throw new QueryCapacityExceededException(config.getMaxConcurrentCalls());
            }
            hallPasses.add(totalLimiter);
        });
        return hallPasses;
    } catch (Exception ex) {
        releaseLanes(hallPasses);
        throw ex;
    }
}

Example 3 with QueryCapacityExceededException

use of org.apache.druid.query.QueryCapacityExceededException in project druid by druid-io.

the class QueryResourceTest method testTooManyQueryInLane.

@Test(timeout = 10_000L)
public void testTooManyQueryInLane() throws InterruptedException {
    expectPermissiveHappyPathAuth();
    final CountDownLatch waitTwoStarted = new CountDownLatch(2);
    final CountDownLatch waitOneScheduled = new CountDownLatch(1);
    final CountDownLatch waitAllFinished = new CountDownLatch(3);
    final QueryScheduler scheduler = new QueryScheduler(40, ManualQueryPrioritizationStrategy.INSTANCE, new HiLoQueryLaningStrategy(2), new ServerConfig());
    createScheduledQueryResource(scheduler, ImmutableList.of(waitTwoStarted), ImmutableList.of(waitOneScheduled));
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY_LOW_PRIORITY, waitAllFinished, response -> Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()));
    waitOneScheduled.await();
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY_LOW_PRIORITY, waitAllFinished, response -> {
        Assert.assertEquals(QueryCapacityExceededException.STATUS_CODE, response.getStatus());
        QueryCapacityExceededException ex;
        try {
            ex = jsonMapper.readValue((byte[]) response.getEntity(), QueryCapacityExceededException.class);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        Assert.assertEquals(QueryCapacityExceededException.makeLaneErrorMessage(HiLoQueryLaningStrategy.LOW, 1), ex.getMessage());
        Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, ex.getErrorCode());
    });
    waitTwoStarted.await();
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY, waitAllFinished, response -> Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()));
    waitAllFinished.await();
}

Example 4 with QueryCapacityExceededException

use of org.apache.druid.query.QueryCapacityExceededException in project druid by druid-io.

the class QueryResourceTest method testTooManyQuery.

@Test(timeout = 10_000L)
public void testTooManyQuery() throws InterruptedException {
    expectPermissiveHappyPathAuth();
    final CountDownLatch waitTwoScheduled = new CountDownLatch(2);
    final CountDownLatch waitAllFinished = new CountDownLatch(3);
    final QueryScheduler laningScheduler = new QueryScheduler(2, ManualQueryPrioritizationStrategy.INSTANCE, NoQueryLaningStrategy.INSTANCE, new ServerConfig());
    createScheduledQueryResource(laningScheduler, Collections.emptyList(), ImmutableList.of(waitTwoScheduled));
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY, waitAllFinished, response -> Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()));
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY, waitAllFinished, response -> Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()));
    waitTwoScheduled.await();
    assertResponseAndCountdownOrBlockForever(SIMPLE_TIMESERIES_QUERY, waitAllFinished, response -> {
        Assert.assertEquals(QueryCapacityExceededException.STATUS_CODE, response.getStatus());
        QueryCapacityExceededException ex;
        try {
            ex = jsonMapper.readValue((byte[]) response.getEntity(), QueryCapacityExceededException.class);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        Assert.assertEquals(QueryCapacityExceededException.makeTotalErrorMessage(2), ex.getMessage());
        Assert.assertEquals(QueryCapacityExceededException.ERROR_CODE, ex.getErrorCode());
    });
    waitAllFinished.await();
}

Example 5 with QueryCapacityExceededException

use of org.apache.druid.query.QueryCapacityExceededException in project druid by druid-io.

the class QueryResource method doPost.

@POST
@Produces({ MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE })
@Consumes({ MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE, APPLICATION_SMILE })
public Response doPost(final InputStream in, @QueryParam("pretty") final String pretty, // used to get request content-type,Accept header, remote address and auth-related headers
@Context final HttpServletRequest req) throws IOException {
    final QueryLifecycle queryLifecycle = queryLifecycleFactory.factorize();
    Query<?> query = null;
    final ResourceIOReaderWriter ioReaderWriter = createResourceIOReaderWriter(req, pretty != null);
    final String currThreadName = Thread.currentThread().getName();
    try {
        queryLifecycle.initialize(readQuery(req, in, ioReaderWriter));
        query = queryLifecycle.getQuery();
        final String queryId = query.getId();
        final String queryThreadName = StringUtils.format("%s[%s_%s_%s]", currThreadName, query.getType(), query.getDataSource().getTableNames(), queryId);
        Thread.currentThread().setName(queryThreadName);
        if (log.isDebugEnabled()) {
            log.debug("Got query [%s]", query);
        }
        final Access authResult = queryLifecycle.authorize(req);
        if (!authResult.isAllowed()) {
            throw new ForbiddenException(authResult.toString());
        }
        final QueryLifecycle.QueryResponse queryResponse = queryLifecycle.execute();
        final Sequence<?> results = queryResponse.getResults();
        final ResponseContext responseContext = queryResponse.getResponseContext();
        final String prevEtag = getPreviousEtag(req);
        if (prevEtag != null && prevEtag.equals(responseContext.getEntityTag())) {
            queryLifecycle.emitLogsAndMetrics(null, req.getRemoteAddr(), -1);
            successfulQueryCount.incrementAndGet();
            return Response.notModified().build();
        }
        final Yielder<?> yielder = Yielders.each(results);
        try {
            boolean shouldFinalize = QueryContexts.isFinalize(query, true);
            boolean serializeDateTimeAsLong = QueryContexts.isSerializeDateTimeAsLong(query, false) || (!shouldFinalize && QueryContexts.isSerializeDateTimeAsLongInner(query, false));
            final ObjectWriter jsonWriter = ioReaderWriter.getResponseWriter().newOutputWriter(queryLifecycle.getToolChest(), queryLifecycle.getQuery(), serializeDateTimeAsLong);
            Response.ResponseBuilder responseBuilder = Response.ok(new StreamingOutput() {

                @Override
                public void write(OutputStream outputStream) throws WebApplicationException {
                    Exception e = null;
                    CountingOutputStream os = new CountingOutputStream(outputStream);
                    try {
                        // json serializer will always close the yielder
                        jsonWriter.writeValue(os, yielder);
                        // Some types of OutputStream suppress flush errors in the .close() method.
                        os.flush();
                        os.close();
                    } catch (Exception ex) {
                        e = ex;
                        log.noStackTrace().error(ex, "Unable to send query response.");
                        throw new RuntimeException(ex);
                    } finally {
                        Thread.currentThread().setName(currThreadName);
                        queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), os.getCount());
                        if (e == null) {
                            successfulQueryCount.incrementAndGet();
                        } else {
                            failedQueryCount.incrementAndGet();
                        }
                    }
                }
            }, ioReaderWriter.getResponseWriter().getResponseType()).header("X-Druid-Query-Id", queryId);
            transferEntityTag(responseContext, responseBuilder);
            DirectDruidClient.removeMagicResponseContextFields(responseContext);
            // Limit the response-context header, see https://github.com/apache/druid/issues/2331
            // Note that Response.ResponseBuilder.header(String key,Object value).build() calls value.toString()
            // and encodes the string using ASCII, so 1 char is = 1 byte
            final ResponseContext.SerializationResult serializationResult = responseContext.serializeWith(jsonMapper, responseContextConfig.getMaxResponseContextHeaderSize());
            if (serializationResult.isTruncated()) {
                final String logToPrint = StringUtils.format("Response Context truncated for id [%s]. Full context is [%s].", queryId, serializationResult.getFullResult());
                if (responseContextConfig.shouldFailOnTruncatedResponseContext()) {
                    log.error(logToPrint);
                    throw new QueryInterruptedException(new TruncatedResponseContextException("Serialized response context exceeds the max size[%s]", responseContextConfig.getMaxResponseContextHeaderSize()), selfNode.getHostAndPortToUse());
                } else {
                    log.warn(logToPrint);
                }
            }
            return responseBuilder.header(HEADER_RESPONSE_CONTEXT, serializationResult.getResult()).build();
        } catch (QueryException e) {
            // make sure to close yielder if anything happened before starting to serialize the response.
            yielder.close();
            throw e;
        } catch (Exception e) {
            // make sure to close yielder if anything happened before starting to serialize the response.
            yielder.close();
            throw new RuntimeException(e);
        } finally {
        // do not close yielder here, since we do not want to close the yielder prior to
        // StreamingOutput having iterated over all the results
        }
    } catch (QueryInterruptedException e) {
        interruptedQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), -1);
        return ioReaderWriter.getResponseWriter().gotError(e);
    } catch (QueryTimeoutException timeout) {
        timedOutQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(timeout, req.getRemoteAddr(), -1);
        return ioReaderWriter.getResponseWriter().gotTimeout(timeout);
    } catch (QueryCapacityExceededException cap) {
        failedQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(cap, req.getRemoteAddr(), -1);
        return ioReaderWriter.getResponseWriter().gotLimited(cap);
    } catch (QueryUnsupportedException unsupported) {
        failedQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(unsupported, req.getRemoteAddr(), -1);
        return ioReaderWriter.getResponseWriter().gotUnsupported(unsupported);
    } catch (BadJsonQueryException | ResourceLimitExceededException e) {
        interruptedQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), -1);
        return ioReaderWriter.getResponseWriter().gotBadQuery(e);
    } catch (ForbiddenException e) {
        // send an error response if this is thrown.
        throw e;
    } catch (Exception e) {
        failedQueryCount.incrementAndGet();
        queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), -1);
        log.noStackTrace().makeAlert(e, "Exception handling request").addData("query", query != null ? jsonMapper.writeValueAsString(query) : "unparseable query").addData("peer", req.getRemoteAddr()).emit();
        return ioReaderWriter.getResponseWriter().gotError(e);
    } finally {
        Thread.currentThread().setName(currThreadName);
    }
}