Examples with Timeout org.junit.rules.Timeout used on opensource projects

Example 41 with Timeout

use of org.junit.rules.Timeout in project pravega by pravega.

the class IndexReaderWriterTests method checkIndex.

private void checkIndex(Collection<BufferView> allKeys, Map<Long, BufferView> existingKeysByOffset, IndexWriter w, KeyHasher hasher, SegmentMock segment) {
    val timer = new TimeoutTimer(TIMEOUT);
    // Group all keys by their full hash (each hash should translate to a bucket), and make sure they're ordered by
    // offset (in descending order - so we can verify backpointer ordering).
    val existingKeys = existingKeysByOffset.entrySet().stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
    val keysByHash = allKeys.stream().map(key -> new BucketUpdate.KeyInfo(key, existingKeys.getOrDefault(key, NO_OFFSET), existingKeys.getOrDefault(key, NO_OFFSET))).sorted(// Reverse order.
    (k1, k2) -> Long.compare(k2.getOffset(), k1.getOffset())).collect(Collectors.groupingBy(keyInfo -> hasher.hash(keyInfo.getKey())));
    int existentBucketCount = 0;
    val buckets = w.locateBuckets(segment, keysByHash.keySet(), timer).join();
    for (val e : keysByHash.entrySet()) {
        val hash = e.getKey();
        val keys = e.getValue();
        val bucket = buckets.get(hash);
        Assert.assertNotNull("No bucket found for hash " + hash, bucket);
        boolean allDeleted = keys.stream().allMatch(k -> k.getOffset() == NO_OFFSET);
        Assert.assertNotEquals("Only expecting inexistent bucket when all its keys are deleted " + hash, allDeleted, bucket.exists());
        val bucketOffsets = w.getBucketOffsets(segment, bucket, timer).join();
        // Verify that we didn't return too many or too few keys.
        if (allDeleted) {
            Assert.assertEquals("Not expecting any offsets to be returned for bucket: " + hash, 0, bucketOffsets.size());
        } else {
            AssertExtensions.assertGreaterThan("Expected at least one offset to be returned for bucket: " + hash, 0, bucketOffsets.size());
            existentBucketCount++;
        }
        AssertExtensions.assertLessThanOrEqual("Too many offsets returned for bucket: " + hash, keys.size(), bucketOffsets.size());
        // Verify returned keys are as expected.
        for (int i = 0; i < bucketOffsets.size(); i++) {
            long actualOffset = bucketOffsets.get(i);
            long expectedOffset = keys.get(i).getOffset();
            String id = String.format("{%s[%s]}", hash, i);
            // In this loop, we do not expect to have Deleted Keys. If our Expected Offset indicates this key should
            // have been deleted, then getBucketOffsets() should not have returned this.
            Assert.assertNotEquals("Expecting a deleted key but found existing one: " + id, NO_OFFSET, expectedOffset);
            Assert.assertEquals("Unexpected key offset in bucket " + id, expectedOffset, actualOffset);
        }
        if (bucketOffsets.size() < keys.size()) {
            val prevKeyOffset = keys.get(bucketOffsets.size()).getOffset();
            Assert.assertEquals("Missing key from bucket " + hash, NO_OFFSET, prevKeyOffset);
        }
    }
    checkEntryCount(existingKeysByOffset.size(), segment);
    checkBucketCount(existentBucketCount, segment);
}

Example 42 with Timeout

use of org.junit.rules.Timeout in project pravega by pravega.

the class ContainerKeyIndexTests method testRecoveryTimeout.

/**
 * Checks the ability for the {@link ContainerKeyIndex} class to cancel recovery-bound tasks if recovery took too long.
 */
@Test
public void testRecoveryTimeout() throws Exception {
    val spyExecutor = Mockito.spy(executorService());
    @Cleanup val context = new TestContext();
    @Cleanup val index = context.createIndex(context.defaultConfig, spyExecutor);
    // Setup the segment with initial attributes.
    val iw = new IndexWriter(HASHER, executorService());
    // Generate initial set of keys.
    val keys = generateUnversionedKeys(BATCH_SIZE, context);
    long offset = 0;
    val hashes = new ArrayList<UUID>();
    val keysWithOffsets = new HashMap<UUID, KeyWithOffset>();
    for (val k : keys) {
        val hash = HASHER.hash(k.getKey());
        hashes.add(hash);
        keysWithOffsets.put(hash, new KeyWithOffset(k.getKey(), offset));
        offset += k.getKey().getLength();
    }
    // Write some garbage data to the segment, but make it longer than the threshold to trigger pre-caching; we don't
    // want to deal with that now since we can't control its runtime.
    context.segment.append(new ByteArraySegment(new byte[TEST_MAX_TAIL_CACHE_PRE_INDEX_LENGTH + 1]), null, TIMEOUT).join();
    // Update the index, but keep the LastIndexedOffset at 0.
    val buckets = iw.locateBuckets(context.segment, keysWithOffsets.keySet(), context.timer).join();
    Collection<BucketUpdate> bucketUpdates = buckets.entrySet().stream().map(e -> {
        val builder = BucketUpdate.forBucket(e.getValue());
        val ko = keysWithOffsets.get(e.getKey());
        builder.withKeyUpdate(new BucketUpdate.KeyUpdate(ko.key, ko.offset, ko.offset, false));
        return builder.build();
    }).collect(Collectors.toList());
    iw.updateBuckets(context.segment, bucketUpdates, 0L, 0, keysWithOffsets.size(), TIMEOUT).join();
    // Setup a mock Executor.schedule call that returns a future over which we have full control. This will save us
    // from having to wait for the recovery timeout to actually expire.
    val timeoutFuture = new AtomicReference<ScheduledFuture<?>>();
    val timeoutCallback = new AtomicReference<Callable<?>>();
    Mockito.doAnswer(arg -> {
        timeoutCallback.set(arg.getArgument(0));
        timeoutFuture.set((ScheduledFuture<?>) arg.callRealMethod());
        return timeoutFuture.get();
    }).when(spyExecutor).schedule(Mockito.<Callable<Boolean>>any(), Mockito.anyLong(), Mockito.any());
    // Issue the first request. We intend to time it out.
    val getBucketOffset1 = index.getBucketOffsets(context.segment, hashes, context.timer);
    TestUtils.await(() -> timeoutFuture.get() != null, 5, TIMEOUT.toMillis());
    // Time-out the call.
    timeoutCallback.get().call();
    AssertExtensions.assertSuppliedFutureThrows("Request did not fail when recovery timed out.", () -> getBucketOffset1, ex -> ex instanceof TimeoutException);
    // Wait for the future to be unregistered before continuing.
    TestUtils.await(timeoutFuture.get()::isDone, 5, TIMEOUT.toMillis());
    // Verify that a new operation will be unblocked if we notify that the recovery completed successfully.
    timeoutFuture.set(null);
    timeoutCallback.set(null);
    val getBucketOffset2 = index.getBucketOffsets(context.segment, hashes, context.timer);
    // Wait for the timeout future to be registered. We'll verify if it gets cancelled next.
    TestUtils.await(() -> timeoutFuture.get() != null, 5, TIMEOUT.toMillis());
    index.notifyIndexOffsetChanged(context.segment.getSegmentId(), context.segment.getInfo().getLength(), 0);
    val result1 = getBucketOffset2.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
    val expected1 = new HashMap<UUID, Long>();
    keysWithOffsets.forEach((k, o) -> expected1.put(k, o.offset));
    AssertExtensions.assertMapEquals("Unexpected result from getBucketOffsets() after a retry.", expected1, result1);
    AssertExtensions.assertEventuallyEquals("Expected the internal timeout future to be cancelled after a successful recovery.", true, timeoutFuture.get()::isCancelled, 5, TIMEOUT.toMillis());
}

Example 43 with Timeout

use of org.junit.rules.Timeout in project pravega by pravega.

the class ContainerKeyIndexTests method checkBackpointers.

private void checkBackpointers(List<UpdateItem> updates, TestContext context) {
    val sortedUpdates = updates.stream().sorted(Comparator.comparingLong(u -> u.offset.get())).collect(Collectors.toList());
    val highestUpdate = sortedUpdates.get(sortedUpdates.size() - 1);
    val highestUpdateHashes = highestUpdate.batch.getItems().stream().map(TableKeyBatch.Item::getHash).collect(Collectors.toList());
    Map<UUID, Long> backpointerSources = context.index.getBucketOffsets(context.segment, highestUpdateHashes, context.timer).join();
    for (int updateId = sortedUpdates.size() - 1; updateId >= 0; updateId--) {
        // Generate the expected backpointers.
        Map<UUID, Long> expectedBackpointers = new HashMap<>();
        if (updateId == 0) {
            // For the first update, we do not expect any.
            backpointerSources.keySet().forEach(k -> expectedBackpointers.put(k, -1L));
        } else {
            // For any other updates, it's whatever got executed before this one.
            val previousUpdate = sortedUpdates.get(updateId - 1);
            for (int i = 0; i < previousUpdate.batch.getItems().size(); i++) {
                expectedBackpointers.put(previousUpdate.batch.getItems().get(i).getHash(), previousUpdate.update.join().get(i));
            }
        }
        // Fetch the actual values.
        Map<UUID, Long> actualBackpointers = backpointerSources.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> context.index.getBackpointerOffset(context.segment, e.getValue(), TIMEOUT).join()));
        // Check and move on.
        Assert.assertEquals("Unexpected backpointer count for update " + updateId, expectedBackpointers.size(), actualBackpointers.size());
        for (val e : expectedBackpointers.entrySet()) {
            val a = actualBackpointers.get(e.getKey());
            Assert.assertNotNull("No backpointer for update " + updateId, a);
            Assert.assertEquals("Unexpected backpointer for update " + updateId, e.getValue(), a);
        }
        backpointerSources = expectedBackpointers;
    }
    // Check unindexed key hashes.
    val unindexedHashes = context.index.getUnindexedKeyHashes(context.segment).join().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().getSegmentOffset()));
    val expectedHashes = highestUpdate.batch.getItems().stream().collect(Collectors.toMap(TableKeyBatch.Item::getHash, i -> highestUpdate.offset.get() + i.getOffset()));
    AssertExtensions.assertMapEquals("Unexpected result from getUnindexedKeyHashes", expectedHashes, unindexedHashes);
}

Example 44 with Timeout

use of org.junit.rules.Timeout in project pravega by pravega.

the class ContainerKeyIndexTests method testGetBucketOffsetDirect.

/**
 * Tests the {@link ContainerKeyIndex#getBucketOffsetDirect} method.
 */
@Test
public void testGetBucketOffsetDirect() {
    final int segmentLength = 1;
    final long updateBatchLength = 100000;
    final long noCacheOffset = updateBatchLength;
    final long lowerCacheOffset = noCacheOffset + updateBatchLength;
    final long higherCacheOffset = lowerCacheOffset + updateBatchLength;
    @Cleanup val context = new TestContext();
    context.segment.append(new ByteArraySegment(new byte[segmentLength]), null, TIMEOUT).join();
    // Setup the segment with initial attributes.
    val iw = new IndexWriter(HASHER, executorService());
    // Generate keys.
    // First 1/3 of the keys do not exist in the cache.
    // Second 1/3 of the keys exist in the cache, but have an offset lower than in the Index.
    // Last 1/3 of the keys exist in the cache and have an offset higher than in the Index.
    val keys = generateUnversionedKeys(BATCH_SIZE, context);
    val keysWithOffsets = new HashMap<UUID, KeyWithOffset>();
    val noCacheKeys = new ArrayList<TableKey>();
    val lowerCacheOffsetKeys = new ArrayList<TableKey>();
    val higherCacheOffsetKeys = new ArrayList<TableKey>();
    for (int i = 0; i < keys.size(); i++) {
        val k = keys.get(i);
        val hash = HASHER.hash(k.getKey());
        if (i < keys.size() / 3) {
            // Does not exist in the cache.
            noCacheKeys.add(k);
            keysWithOffsets.put(hash, new KeyWithOffset(k.getKey(), noCacheOffset));
        } else if (i < keys.size() * 2 / 3) {
            // Exists in the cache, but with a lower offset than in the index.
            lowerCacheOffsetKeys.add(k);
            keysWithOffsets.put(hash, new KeyWithOffset(k.getKey(), lowerCacheOffset));
        } else {
            // Exists in the cache with a higher offset than in the index.
            higherCacheOffsetKeys.add(k);
            keysWithOffsets.put(hash, new KeyWithOffset(k.getKey(), higherCacheOffset));
        }
    }
    // Update everything in the underlying index.
    val buckets = iw.locateBuckets(context.segment, keysWithOffsets.keySet(), context.timer).join();
    Collection<BucketUpdate> bucketUpdates = buckets.entrySet().stream().map(e -> {
        val builder = BucketUpdate.forBucket(e.getValue());
        val ko = keysWithOffsets.get(e.getKey());
        builder.withKeyUpdate(new BucketUpdate.KeyUpdate(ko.key, ko.offset, ko.offset, false));
        return builder.build();
    }).collect(Collectors.toList());
    iw.updateBuckets(context.segment, bucketUpdates, 0L, segmentLength, 0, TIMEOUT).join();
    // Update cache, and immediately clear out the tail section as we want to simulate a case where the values are already
    // thought to be indexed already.
    context.index.update(context.segment, toUpdateBatch(lowerCacheOffsetKeys), () -> CompletableFuture.completedFuture(lowerCacheOffset - updateBatchLength), context.timer).join();
    context.index.update(context.segment, toUpdateBatch(higherCacheOffsetKeys), () -> CompletableFuture.completedFuture(higherCacheOffset + BATCH_SIZE), context.timer).join();
    context.index.notifyIndexOffsetChanged(context.segment.getSegmentId(), higherCacheOffset + updateBatchLength, 0);
    // Check results. The expected offsets should already be stored in keysWithOffsets.
    for (val k : keys) {
        val hash = HASHER.hash(k.getKey());
        val actualOffset = context.index.getBucketOffsetDirect(context.segment, hash, context.timer).join();
        val expectedOffset = keysWithOffsets.get(hash).offset;
        Assert.assertEquals("Unexpected result from getBucketOffsetDirect.", expectedOffset, (long) actualOffset);
        val cachedOffset = context.index.getBucketOffsets(context.segment, Collections.singleton(hash), context.timer).join().get(hash);
        Assert.assertEquals("Unexpected result from getBucketOffsets.", expectedOffset, (long) cachedOffset);
    }
}

Example 45 with Timeout

use of org.junit.rules.Timeout in project pravega by pravega.

the class WriterTableProcessorTests method getRemovalLengths.

private long getRemovalLengths(long truncationOffset, List<TestBatchData> batches, TestContext context) throws Exception {
    val candidates = batches.stream().flatMap(b -> b.operations.stream()).filter(op -> op.getStreamSegmentOffset() >= truncationOffset).iterator();
    val expectedEntries = batches.get(batches.size() - 1).expectedEntries;
    long result = 0;
    while (candidates.hasNext()) {
        val op = candidates.next();
        val opData = new byte[(int) op.getLength()];
        val bytesRead = context.segmentMock.read(op.getStreamSegmentOffset(), (int) op.getLength(), TIMEOUT).readRemaining(opData, TIMEOUT);
        assert bytesRead == opData.length;
        val entryHeader = context.serializer.readHeader(new ByteArraySegment(opData).getBufferViewReader());
        if (!expectedEntries.containsKey(new ByteArraySegment(opData, entryHeader.getKeyOffset(), entryHeader.getKeyLength()))) {
            result += op.getLength();
        } else {
            break;
        }
    }
    return result;
}