Examples with OutgoingStreamMessage org.apache.cassandra.streaming.messages.OutgoingStreamMessage used on opensource projects

Example 6 with OutgoingStreamMessage

use of org.apache.cassandra.streaming.messages.OutgoingStreamMessage in project cassandra by apache.

the class StreamTransferTaskTest method testFailSessionDuringTransferShouldNotReleaseReferences.

@Test
public void testFailSessionDuringTransferShouldNotReleaseReferences() throws Exception {
    InetAddressAndPort peer = FBUtilities.getBroadcastAddressAndPort();
    StreamCoordinator streamCoordinator = new StreamCoordinator(StreamOperation.BOOTSTRAP, 1, new NettyStreamingConnectionFactory(), false, false, null, PreviewKind.NONE);
    StreamResultFuture future = StreamResultFuture.createInitiator(UUID.randomUUID(), StreamOperation.OTHER, Collections.<StreamEventHandler>emptyList(), streamCoordinator);
    StreamSession session = new StreamSession(StreamOperation.BOOTSTRAP, peer, FACTORY, null, current_version, false, 0, null, PreviewKind.NONE);
    session.init(future);
    ColumnFamilyStore cfs = Keyspace.open(KEYSPACE1).getColumnFamilyStore(CF_STANDARD);
    // create two sstables
    for (int i = 0; i < 2; i++) {
        SchemaLoader.insertData(KEYSPACE1, CF_STANDARD, i, 1);
        cfs.forceBlockingFlush();
    }
    // create streaming task that streams those two sstables
    StreamTransferTask task = new StreamTransferTask(session, cfs.metadata.id);
    List<Ref<SSTableReader>> refs = new ArrayList<>(cfs.getLiveSSTables().size());
    for (SSTableReader sstable : cfs.getLiveSSTables()) {
        List<Range<Token>> ranges = new ArrayList<>();
        ranges.add(new Range<>(sstable.first.getToken(), sstable.last.getToken()));
        Ref<SSTableReader> ref = sstable.selfRef();
        refs.add(ref);
        task.addTransferStream(new CassandraOutgoingFile(StreamOperation.BOOTSTRAP, ref, sstable.getPositionsForRanges(ranges), ranges, 1));
    }
    assertEquals(14, task.getTotalNumberOfFiles());
    // add task to stream session, so it is aborted when stream session fails
    session.transfers.put(TableId.generate(), task);
    // make a copy of outgoing file messages, since task is cleared when it's aborted
    Collection<OutgoingStreamMessage> files = new LinkedList<>(task.streams.values());
    // simulate start transfer
    for (OutgoingStreamMessage file : files) {
        file.startTransfer();
    }
    // fail stream session mid-transfer
    session.onError(new Exception("Fake exception")).get(5, TimeUnit.SECONDS);
    // make sure reference was not released
    for (Ref<SSTableReader> ref : refs) {
        assertEquals(1, ref.globalCount());
    }
    // wait for stream to abort asynchronously
    int tries = 10;
    while (ScheduledExecutors.nonPeriodicTasks.getActiveTaskCount() > 0) {
        if (tries < 1)
            throw new RuntimeException("test did not complete in time");
        Thread.sleep(10);
        tries--;
    }
    // simulate finish transfer
    for (OutgoingStreamMessage file : files) {
        file.finishTransfer();
    }
    // now reference should be released
    for (Ref<SSTableReader> ref : refs) {
        assertEquals(0, ref.globalCount());
    }
}