Examples with ServerSegmentCompletionProtocolHandler com.linkedin.pinot.server.realtime.ServerSegmentCompletionProtocolHandler used on opensource projects

Example 1 with ServerSegmentCompletionProtocolHandler

use of com.linkedin.pinot.server.realtime.ServerSegmentCompletionProtocolHandler in project pinot by linkedin.

the class SegmentCompletionIntegrationTests method testStopConsumingToOfflineAndAutofix.

// Test that if we send stoppedConsuming to the controller, the segment goes offline.
@Test
public void testStopConsumingToOfflineAndAutofix() throws Exception {
    final String realtimeTableName = TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(_tableName);
    long endTime = now() + MAX_RUN_TIME_SECONDS * 1000L;
    while (now() < endTime) {
        ExternalView ev = HelixHelper.getExternalViewForResource(_helixAdmin, _clusterName, realtimeTableName);
        if (ev != null) {
            Map<String, String> stateMap = ev.getStateMap(_segmentName);
            if (stateMap != null && stateMap.containsKey(_serverInstance)) {
                if (stateMap.get(_serverInstance).equals(PinotHelixSegmentOnlineOfflineStateModelGenerator.CONSUMING_STATE)) {
                    break;
                }
            }
        }
        Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
    }
    Assert.assertTrue(now() < endTime, "Failed trying to reach consuming state");
    // Now report to the controller that we had to stop consumption
    ServerSegmentCompletionProtocolHandler protocolHandler = new ServerSegmentCompletionProtocolHandler(_serverInstance);
    SegmentCompletionProtocol.Request.Params params = new SegmentCompletionProtocol.Request.Params();
    params.withOffset(45688L).withSegmentName(_segmentName).withReason("RandomReason");
    SegmentCompletionProtocol.Response response = protocolHandler.segmentStoppedConsuming(params);
    Assert.assertTrue(response.getStatus() == SegmentCompletionProtocol.ControllerResponseStatus.PROCESSED);
    while (now() < endTime) {
        ExternalView ev = HelixHelper.getExternalViewForResource(_helixAdmin, _clusterName, realtimeTableName);
        Map<String, String> stateMap = ev.getStateMap(_segmentName);
        if (stateMap.containsKey(_serverInstance)) {
            if (stateMap.get(_serverInstance).equals(PinotHelixSegmentOnlineOfflineStateModelGenerator.OFFLINE_STATE)) {
                break;
            }
        }
        Uninterruptibles.sleepUninterruptibly(50, TimeUnit.MILLISECONDS);
    }
    Assert.assertTrue(now() < endTime, "Failed trying to reach offline state");
    // Now call the validation manager, and the segment should fix itself
    getControllerValidationManager().runValidation();
    // Now there should be a new segment in CONSUMING state in the IDEALSTATE.
    IdealState idealState = HelixHelper.getTableIdealState(_helixManager, realtimeTableName);
    Assert.assertEquals(idealState.getPartitionSet().size(), 2);
    for (String segmentId : idealState.getPartitionSet()) {
        if (!segmentId.equals(_segmentName)) {
            // This is a new segment. Verify that it is in CONSUMING state, and has a sequence number 1 more than prev one
            LLCSegmentName oldSegmentName = new LLCSegmentName(_segmentName);
            LLCSegmentName newSegmentName = new LLCSegmentName(segmentId);
            Assert.assertEquals(newSegmentName.getSequenceNumber(), oldSegmentName.getSequenceNumber() + 1);
            Map<String, String> instanceStateMap = idealState.getInstanceStateMap(segmentId);
            for (String state : instanceStateMap.values()) {
                Assert.assertTrue(state.equals(PinotHelixSegmentOnlineOfflineStateModelGenerator.CONSUMING_STATE));
            }
        }
    }
// We will assume that it eventually makes it to externalview
}