Examples with SequenceWrapper org.apache.druid.java.util.common.guava.SequenceWrapper used on opensource projects

Example 1 with SequenceWrapper

use of org.apache.druid.java.util.common.guava.SequenceWrapper in project druid by druid-io.

the class ResultLevelCachingQueryRunner method run.

@Override
public Sequence<T> run(QueryPlus queryPlus, ResponseContext responseContext) {
    if (useResultCache || populateResultCache) {
        final String cacheKeyStr = StringUtils.fromUtf8(strategy.computeResultLevelCacheKey(query));
        final byte[] cachedResultSet = fetchResultsFromResultLevelCache(cacheKeyStr);
        String existingResultSetId = extractEtagFromResults(cachedResultSet);
        existingResultSetId = existingResultSetId == null ? "" : existingResultSetId;
        query = query.withOverriddenContext(ImmutableMap.of(QueryResource.HEADER_IF_NONE_MATCH, existingResultSetId));
        Sequence<T> resultFromClient = baseRunner.run(QueryPlus.wrap(query), responseContext);
        String newResultSetId = responseContext.getEntityTag();
        if (useResultCache && newResultSetId != null && newResultSetId.equals(existingResultSetId)) {
            log.debug("Return cached result set as there is no change in identifiers for query %s ", query.getId());
            return deserializeResults(cachedResultSet, strategy, existingResultSetId);
        } else {
            @Nullable ResultLevelCachePopulator resultLevelCachePopulator = createResultLevelCachePopulator(cacheKeyStr, newResultSetId);
            if (resultLevelCachePopulator == null) {
                return resultFromClient;
            }
            final Function<T, Object> cacheFn = strategy.prepareForCache(true);
            return Sequences.wrap(Sequences.map(resultFromClient, new Function<T, T>() {

                @Override
                public T apply(T input) {
                    if (resultLevelCachePopulator.isShouldPopulate()) {
                        resultLevelCachePopulator.cacheResultEntry(input, cacheFn);
                    }
                    return input;
                }
            }), new SequenceWrapper() {

                @Override
                public void after(boolean isDone, Throwable thrown) {
                    Preconditions.checkNotNull(resultLevelCachePopulator, "ResultLevelCachePopulator cannot be null during cache population");
                    if (thrown != null) {
                        log.error(thrown, "Error while preparing for result level caching for query %s with error %s ", query.getId(), thrown.getMessage());
                    } else if (resultLevelCachePopulator.isShouldPopulate()) {
                        // The resultset identifier and its length is cached along with the resultset
                        resultLevelCachePopulator.populateResults();
                        log.debug("Cache population complete for query %s", query.getId());
                    }
                    resultLevelCachePopulator.stopPopulating();
                }
            });
        }
    } else {
        return baseRunner.run(queryPlus, responseContext);
    }
}

Example 2 with SequenceWrapper

use of org.apache.druid.java.util.common.guava.SequenceWrapper in project druid by druid-io.

the class QueryLifecycle method runSimple.

/**
 * For callers who have already authorized their query, and where simplicity is desired over flexibility. This method
 * does it all in one call. Logs and metrics are emitted when the Sequence is either fully iterated or throws an
 * exception.
 *
 * @param query                 the query
 * @param authenticationResult  authentication result indicating identity of the requester
 * @param authorizationResult   authorization result of requester
 *
 * @return results
 */
@SuppressWarnings("unchecked")
public <T> Sequence<T> runSimple(final Query<T> query, final AuthenticationResult authenticationResult, final Access authorizationResult) {
    initialize(query);
    final Sequence<T> results;
    try {
        preAuthorized(authenticationResult, authorizationResult);
        if (!authorizationResult.isAllowed()) {
            throw new ISE("Unauthorized");
        }
        final QueryLifecycle.QueryResponse queryResponse = execute();
        results = queryResponse.getResults();
    } catch (Throwable e) {
        emitLogsAndMetrics(e, null, -1);
        throw e;
    }
    return Sequences.wrap(results, new SequenceWrapper() {

        @Override
        public void after(final boolean isDone, final Throwable thrown) {
            emitLogsAndMetrics(thrown, null, -1);
        }
    });
}

Example 3 with SequenceWrapper

use of org.apache.druid.java.util.common.guava.SequenceWrapper in project druid by druid-io.

the class CachingQueryRunnerTest method testCloseAndPopulate.

private void testCloseAndPopulate(List<Result> expectedRes, List<Result> expectedCacheRes, Query query, QueryToolChest toolchest) throws Exception {
    final AssertingClosable closable = new AssertingClosable();
    final Sequence resultSeq = Sequences.wrap(Sequences.simple(expectedRes), new SequenceWrapper() {

        @Override
        public void before() {
            Assert.assertFalse(closable.isClosed());
        }

        @Override
        public void after(boolean isDone, Throwable thrown) {
            closable.close();
        }
    });
    final CountDownLatch cacheMustBePutOnce = new CountDownLatch(1);
    Cache cache = new Cache() {

        private final ConcurrentMap<NamedKey, byte[]> baseMap = new ConcurrentHashMap<>();

        @Override
        public byte[] get(NamedKey key) {
            return baseMap.get(key);
        }

        @Override
        public void put(NamedKey key, byte[] value) {
            baseMap.put(key, value);
            cacheMustBePutOnce.countDown();
        }

        @Override
        public Map<NamedKey, byte[]> getBulk(Iterable<NamedKey> keys) {
            return null;
        }

        @Override
        public void close(String namespace) {
        }

        @Override
        public void close() {
        }

        @Override
        public CacheStats getStats() {
            return null;
        }

        @Override
        public boolean isLocal() {
            return true;
        }

        @Override
        public void doMonitor(ServiceEmitter emitter) {
        }
    };
    byte[] keyPrefix = RandomUtils.nextBytes(10);
    CachingQueryRunner runner = makeCachingQueryRunner(keyPrefix, cache, toolchest, resultSeq);
    CacheStrategy cacheStrategy = toolchest.getCacheStrategy(query);
    Cache.NamedKey cacheKey = CacheUtil.computeSegmentCacheKey(CACHE_ID, SEGMENT_DESCRIPTOR, Bytes.concat(keyPrefix, cacheStrategy.computeCacheKey(query)));
    Assert.assertTrue(runner.canPopulateCache(query, cacheStrategy));
    Sequence res = runner.run(QueryPlus.wrap(query));
    // base sequence is not closed yet
    Assert.assertFalse("sequence must not be closed", closable.isClosed());
    Assert.assertNull("cache must be empty", cache.get(cacheKey));
    List results = res.toList();
    Assert.assertTrue(closable.isClosed());
    Assert.assertEquals(expectedRes.toString(), results.toString());
    // wait for background caching finish
    // wait at most 10 seconds to fail the test to avoid block overall tests
    Assert.assertTrue("cache must be populated", cacheMustBePutOnce.await(10, TimeUnit.SECONDS));
    byte[] cacheValue = cache.get(cacheKey);
    Assert.assertNotNull(cacheValue);
    Function<Object, Result> fn = cacheStrategy.pullFromSegmentLevelCache();
    List<Result> cacheResults = Lists.newArrayList(Iterators.transform(objectMapper.readValues(objectMapper.getFactory().createParser(cacheValue), cacheStrategy.getCacheObjectClazz()), fn));
    Assert.assertEquals(expectedCacheRes.toString(), cacheResults.toString());
}

Example 4 with SequenceWrapper

use of org.apache.druid.java.util.common.guava.SequenceWrapper in project druid by druid-io.

the class QuerySchedulerTest method testHiLoReleaseLaneWhenSequenceExplodes.

@Test
public void testHiLoReleaseLaneWhenSequenceExplodes() throws Exception {
    expected.expectMessage("exploded");
    expected.expect(ExecutionException.class);
    TopNQuery interactive = makeInteractiveQuery();
    ListenableFuture<?> future = executorService.submit(() -> {
        try {
            Query<?> scheduled = scheduler.prioritizeAndLaneQuery(QueryPlus.wrap(interactive), ImmutableSet.of());
            Assert.assertNotNull(scheduled);
            Sequence<Integer> underlyingSequence = makeExplodingSequence(10);
            underlyingSequence = Sequences.wrap(underlyingSequence, new SequenceWrapper() {

                @Override
                public void before() {
                    Assert.assertEquals(4, scheduler.getTotalAvailableCapacity());
                }
            });
            Sequence<Integer> results = scheduler.run(scheduled, underlyingSequence);
            consumeAndCloseSequence(results);
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
    });
    future.get();
}

Example 5 with SequenceWrapper

use of org.apache.druid.java.util.common.guava.SequenceWrapper in project druid by druid-io.

the class SqlLifecycle method runSimple.

@VisibleForTesting
public Sequence<Object[]> runSimple(String sql, Map<String, Object> queryContext, List<SqlParameter> parameters, AuthenticationResult authenticationResult) throws RelConversionException {
    Sequence<Object[]> result;
    initialize(sql, queryContext);
    try {
        setParameters(SqlQuery.getParameterList(parameters));
        validateAndAuthorize(authenticationResult);
        plan();
        result = execute();
    } catch (Throwable e) {
        if (!(e instanceof ForbiddenException)) {
            finalizeStateAndEmitLogsAndMetrics(e, null, -1);
        }
        throw e;
    }
    return Sequences.wrap(result, new SequenceWrapper() {

        @Override
        public void after(boolean isDone, Throwable thrown) {
            finalizeStateAndEmitLogsAndMetrics(thrown, null, -1);
        }
    });
}