Examples with ReferenceCountingSegment org.apache.druid.segment.ReferenceCountingSegment used on opensource projects

Example 26 with ReferenceCountingSegment

use of org.apache.druid.segment.ReferenceCountingSegment in project druid by druid-io.

the class ServerManagerForQueryErrorTest method buildQueryRunnerForSegment.

@Override
protected <T> QueryRunner<T> buildQueryRunnerForSegment(Query<T> query, SegmentDescriptor descriptor, QueryRunnerFactory<T, Query<T>> factory, QueryToolChest<T, Query<T>> toolChest, VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline, Function<SegmentReference, SegmentReference> segmentMapFn, AtomicLong cpuTimeAccumulator, Optional<byte[]> cacheKeyPrefix) {
    if (query.getContextBoolean(QUERY_RETRY_TEST_CONTEXT_KEY, false)) {
        final MutableBoolean isIgnoreSegment = new MutableBoolean(false);
        queryToIgnoredSegments.compute(query.getMostSpecificId(), (queryId, ignoredSegments) -> {
            if (ignoredSegments == null) {
                ignoredSegments = new HashSet<>();
            }
            if (ignoredSegments.size() < MAX_NUM_FALSE_MISSING_SEGMENTS_REPORTS) {
                ignoredSegments.add(descriptor);
                isIgnoreSegment.setTrue();
            }
            return ignoredSegments;
        });
        if (isIgnoreSegment.isTrue()) {
            LOG.info("Pretending I don't have segment[%s]", descriptor);
            return new ReportTimelineMissingSegmentQueryRunner<>(descriptor);
        }
    } else if (query.getContextBoolean(QUERY_TIMEOUT_TEST_CONTEXT_KEY, false)) {
        return (queryPlus, responseContext) -> new Sequence<T>() {

            @Override
            public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator) {
                throw new QueryTimeoutException("query timeout test");
            }

            @Override
            public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulator<OutType, T> accumulator) {
                throw new QueryTimeoutException("query timeout test");
            }
        };
    } else if (query.getContextBoolean(QUERY_CAPACITY_EXCEEDED_TEST_CONTEXT_KEY, false)) {
        return (queryPlus, responseContext) -> new Sequence<T>() {

            @Override
            public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator) {
                throw QueryCapacityExceededException.withErrorMessageAndResolvedHost("query capacity exceeded test");
            }

            @Override
            public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulator<OutType, T> accumulator) {
                throw QueryCapacityExceededException.withErrorMessageAndResolvedHost("query capacity exceeded test");
            }
        };
    } else if (query.getContextBoolean(QUERY_UNSUPPORTED_TEST_CONTEXT_KEY, false)) {
        return (queryPlus, responseContext) -> new Sequence<T>() {

            @Override
            public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator) {
                throw new QueryUnsupportedException("query unsupported test");
            }

            @Override
            public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulator<OutType, T> accumulator) {
                throw new QueryUnsupportedException("query unsupported test");
            }
        };
    } else if (query.getContextBoolean(RESOURCE_LIMIT_EXCEEDED_TEST_CONTEXT_KEY, false)) {
        return (queryPlus, responseContext) -> new Sequence<T>() {

            @Override
            public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator) {
                throw new ResourceLimitExceededException("resource limit exceeded test");
            }

            @Override
            public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulator<OutType, T> accumulator) {
                throw new ResourceLimitExceededException("resource limit exceeded test");
            }
        };
    } else if (query.getContextBoolean(QUERY_FAILURE_TEST_CONTEXT_KEY, false)) {
        return (queryPlus, responseContext) -> new Sequence<T>() {

            @Override
            public <OutType> OutType accumulate(OutType initValue, Accumulator<OutType, T> accumulator) {
                throw new RuntimeException("query failure test");
            }

            @Override
            public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulator<OutType, T> accumulator) {
                throw new RuntimeException("query failure test");
            }
        };
    }
    return super.buildQueryRunnerForSegment(query, descriptor, factory, toolChest, timeline, segmentMapFn, cpuTimeAccumulator, cacheKeyPrefix);
}

Example 27 with ReferenceCountingSegment

use of org.apache.druid.segment.ReferenceCountingSegment in project druid by druid-io.

the class SegmentManagerTest method assertResult.

// TODO remove when the bug in intelliJ is fixed.
@SuppressWarnings("RedundantThrows")
private void assertResult(List<DataSegment> expectedExistingSegments) throws SegmentLoadingException {
    final Map<String, Long> expectedDataSourceSizes = expectedExistingSegments.stream().collect(Collectors.toMap(DataSegment::getDataSource, DataSegment::getSize, Long::sum));
    final Map<String, Long> expectedDataSourceCounts = expectedExistingSegments.stream().collect(Collectors.toMap(DataSegment::getDataSource, segment -> 1L, Long::sum));
    final Set<String> expectedDataSourceNames = expectedExistingSegments.stream().map(DataSegment::getDataSource).collect(Collectors.toSet());
    final Map<String, VersionedIntervalTimeline<String, ReferenceCountingSegment>> expectedTimelines = new HashMap<>();
    for (DataSegment segment : expectedExistingSegments) {
        final VersionedIntervalTimeline<String, ReferenceCountingSegment> expectedTimeline = expectedTimelines.computeIfAbsent(segment.getDataSource(), k -> new VersionedIntervalTimeline<>(Ordering.natural()));
        expectedTimeline.add(segment.getInterval(), segment.getVersion(), segment.getShardSpec().createChunk(ReferenceCountingSegment.wrapSegment(SEGMENT_LOADER.getSegment(segment, false, SegmentLazyLoadFailCallback.NOOP), segment.getShardSpec())));
    }
    Assert.assertEquals(expectedDataSourceNames, segmentManager.getDataSourceNames());
    Assert.assertEquals(expectedDataSourceCounts, segmentManager.getDataSourceCounts());
    Assert.assertEquals(expectedDataSourceSizes, segmentManager.getDataSourceSizes());
    final Map<String, DataSourceState> dataSources = segmentManager.getDataSources();
    Assert.assertEquals(expectedTimelines.size(), dataSources.size());
    dataSources.forEach((sourceName, dataSourceState) -> {
        Assert.assertEquals(expectedDataSourceCounts.get(sourceName).longValue(), dataSourceState.getNumSegments());
        Assert.assertEquals(expectedDataSourceSizes.get(sourceName).longValue(), dataSourceState.getTotalSegmentSize());
        Assert.assertEquals(expectedTimelines.get(sourceName).getAllTimelineEntries(), dataSourceState.getTimeline().getAllTimelineEntries());
    });
}

Example 28 with ReferenceCountingSegment

use of org.apache.druid.segment.ReferenceCountingSegment in project druid by druid-io.

the class ServerManagerTest method testMultipleDrops.

@Test
public void testMultipleDrops() throws Exception {
    loadQueryable("test", "3", Intervals.of("2011-04-04/2011-04-05"));
    Future future = assertQueryable(Granularities.DAY, "test", Intervals.of("2011-04-04/2011-04-06"), ImmutableList.of(new Pair<String, Interval>("3", Intervals.of("2011-04-04/2011-04-05"))));
    queryNotifyLatch.await(1000, TimeUnit.MILLISECONDS);
    Assert.assertEquals(1, factory.getSegmentReferences().size());
    for (ReferenceCountingSegment referenceCountingSegment : factory.getSegmentReferences()) {
        Assert.assertEquals(1, referenceCountingSegment.getNumReferences());
    }
    queryWaitYieldLatch.countDown();
    Assert.assertEquals(1, factory.getAdapters().size());
    for (SegmentForTesting segmentForTesting : factory.getAdapters()) {
        Assert.assertFalse(segmentForTesting.isClosed());
    }
    dropQueryable("test", "3", Intervals.of("2011-04-04/2011-04-05"));
    dropQueryable("test", "3", Intervals.of("2011-04-04/2011-04-05"));
    for (SegmentForTesting segmentForTesting : factory.getAdapters()) {
        Assert.assertFalse(segmentForTesting.isClosed());
    }
    queryWaitLatch.countDown();
    future.get();
    for (SegmentForTesting segmentForTesting : factory.getAdapters()) {
        Assert.assertTrue(segmentForTesting.isClosed());
    }
}

Example 29 with ReferenceCountingSegment

use of org.apache.druid.segment.ReferenceCountingSegment in project druid by druid-io.

the class ServerManagerTest method testGetQueryRunnerForSegmentsWhenSegmentIsClosedReportingMissingSegments.

@Test
public void testGetQueryRunnerForSegmentsWhenSegmentIsClosedReportingMissingSegments() {
    final Interval interval = Intervals.of("P1d/2011-04-01");
    final SearchQuery query = searchQuery("test", interval, Granularities.ALL);
    final Optional<VersionedIntervalTimeline<String, ReferenceCountingSegment>> maybeTimeline = segmentManager.getTimeline(DataSourceAnalysis.forDataSource(query.getDataSource()));
    Assert.assertTrue(maybeTimeline.isPresent());
    final List<TimelineObjectHolder<String, ReferenceCountingSegment>> holders = maybeTimeline.get().lookup(interval);
    final List<SegmentDescriptor> closedSegments = new ArrayList<>();
    for (TimelineObjectHolder<String, ReferenceCountingSegment> holder : holders) {
        for (PartitionChunk<ReferenceCountingSegment> chunk : holder.getObject()) {
            final ReferenceCountingSegment segment = chunk.getObject();
            Assert.assertNotNull(segment.getId());
            closedSegments.add(new SegmentDescriptor(segment.getDataInterval(), segment.getVersion(), segment.getId().getPartitionNum()));
            segment.close();
        }
    }
    final QueryRunner<Result<SearchResultValue>> queryRunner = serverManager.getQueryRunnerForSegments(query, closedSegments);
    final ResponseContext responseContext = DefaultResponseContext.createEmpty();
    final List<Result<SearchResultValue>> results = queryRunner.run(QueryPlus.wrap(query), responseContext).toList();
    Assert.assertTrue(results.isEmpty());
    Assert.assertNotNull(responseContext.getMissingSegments());
    Assert.assertEquals(closedSegments, responseContext.getMissingSegments());
}

Example 30 with ReferenceCountingSegment

use of org.apache.druid.segment.ReferenceCountingSegment in project druid by druid-io.

the class FireHydrantTest method testGetIncrementedSegmentClosed.

@Test
public void testGetIncrementedSegmentClosed() {
    expectedException.expect(ISE.class);
    expectedException.expectMessage("segment.close() is called somewhere outside FireHydrant.swapSegment()");
    hydrant.getHydrantSegment().close();
    Assert.assertEquals(0, hydrant.getHydrantSegment().getNumReferences());
    ReferenceCountingSegment segment = hydrant.getIncrementedSegment();
}