Examples with TestReadResultHandler io.pravega.segmentstore.server.reading.TestReadResultHandler used on opensource projects

Example 1 with TestReadResultHandler

use of io.pravega.segmentstore.server.reading.TestReadResultHandler in project pravega by pravega.

the class StreamSegmentContainerTests method testFutureReads.

/**
 * Tests the ability to perform future (tail) reads. Scenarios tested include:
 * * Regular appends
 * * Segment sealing
 * * Transaction merging.
 */
@Test
public void testFutureReads() throws Exception {
    final int nonSealReadLimit = 100;
    @Cleanup TestContext context = new TestContext();
    context.container.startAsync().awaitRunning();
    // 1. Create the StreamSegments.
    ArrayList<String> segmentNames = createSegments(context);
    HashMap<String, ArrayList<String>> transactionsBySegment = createTransactions(segmentNames, context);
    activateAllSegments(segmentNames, context);
    transactionsBySegment.values().forEach(s -> activateAllSegments(s, context));
    HashMap<String, ReadResult> readsBySegment = new HashMap<>();
    ArrayList<AsyncReadResultProcessor> readProcessors = new ArrayList<>();
    HashSet<String> segmentsToSeal = new HashSet<>();
    HashMap<String, ByteArrayOutputStream> readContents = new HashMap<>();
    HashMap<String, TestReadResultHandler> entryHandlers = new HashMap<>();
    // should stop upon reaching the limit).
    for (int i = 0; i < segmentNames.size(); i++) {
        String segmentName = segmentNames.get(i);
        ByteArrayOutputStream readContentsStream = new ByteArrayOutputStream();
        readContents.put(segmentName, readContentsStream);
        ReadResult readResult;
        if (i < segmentNames.size() / 2) {
            // We're going to seal this one at one point.
            segmentsToSeal.add(segmentName);
            readResult = context.container.read(segmentName, 0, Integer.MAX_VALUE, TIMEOUT).join();
        } else {
            // Just a regular one, nothing special.
            readResult = context.container.read(segmentName, 0, nonSealReadLimit, TIMEOUT).join();
        }
        // The Read callback is only accumulating data in this test; we will then compare it against the real data.
        TestReadResultHandler entryHandler = new TestReadResultHandler(readContentsStream, TIMEOUT);
        entryHandlers.put(segmentName, entryHandler);
        readsBySegment.put(segmentName, readResult);
        readProcessors.add(AsyncReadResultProcessor.process(readResult, entryHandler, executorService()));
    }
    // 3. Add some appends.
    HashMap<String, Long> lengths = new HashMap<>();
    HashMap<String, ByteArrayOutputStream> segmentContents = new HashMap<>();
    appendToParentsAndTransactions(segmentNames, transactionsBySegment, lengths, segmentContents, context);
    // 4. Merge all the Transactions.
    mergeTransactions(transactionsBySegment, lengths, segmentContents, context);
    // 5. Add more appends (to the parent segments)
    ArrayList<CompletableFuture<Void>> operationFutures = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        for (String segmentName : segmentNames) {
            byte[] appendData = getAppendData(segmentName, APPENDS_PER_SEGMENT + i);
            operationFutures.add(context.container.append(segmentName, appendData, null, TIMEOUT));
            lengths.put(segmentName, lengths.getOrDefault(segmentName, 0L) + appendData.length);
            recordAppend(segmentName, appendData, segmentContents);
        }
    }
    segmentsToSeal.forEach(segmentName -> operationFutures.add(Futures.toVoid(context.container.sealStreamSegment(segmentName, TIMEOUT))));
    Futures.allOf(operationFutures).get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
    // Now wait for all the reads to complete, and verify their results against the expected output.
    Futures.allOf(entryHandlers.values().stream().map(h -> h.getCompleted()).collect(Collectors.toList())).get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
    readProcessors.forEach(AsyncReadResultProcessor::close);
    // Check to see if any errors got thrown (and caught) during the reading process).
    for (Map.Entry<String, TestReadResultHandler> e : entryHandlers.entrySet()) {
        Throwable err = e.getValue().getError().get();
        if (err != null) {
            // The next check (see below) will verify if the segments were properly read).
            if (!(err instanceof StreamSegmentSealedException && segmentsToSeal.contains(e.getKey()))) {
                Assert.fail("Unexpected error happened while processing Segment " + e.getKey() + ": " + e.getValue().getError().get());
            }
        }
    }
    // Check that all the ReadResults are closed
    for (Map.Entry<String, ReadResult> e : readsBySegment.entrySet()) {
        Assert.assertTrue("Read result is not closed for segment " + e.getKey(), e.getValue().isClosed());
    }
    // Compare, byte-by-byte, the outcome of the tail reads.
    Assert.assertEquals("Unexpected number of segments were read.", segmentContents.size(), readContents.size());
    for (String segmentName : segmentNames) {
        boolean isSealed = segmentsToSeal.contains(segmentName);
        byte[] expectedData = segmentContents.get(segmentName).toByteArray();
        byte[] actualData = readContents.get(segmentName).toByteArray();
        int expectedLength = isSealed ? (int) (long) lengths.get(segmentName) : nonSealReadLimit;
        Assert.assertEquals("Unexpected read length for segment " + segmentName, expectedLength, actualData.length);
        AssertExtensions.assertArrayEquals("Unexpected read contents for segment " + segmentName, expectedData, 0, actualData, 0, actualData.length);
    }
    // 6. Writer moving data to Storage.
    waitForSegmentsInStorage(segmentNames, context).get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
    checkStorage(segmentContents, lengths, context);
}