Search in sources :

Example 16 with Criteria

use of org.apache.helix.Criteria in project incubator-gobblin by apache.

the class YarnAppSecurityManager method sendTokenFileUpdatedMessage.

@VisibleForTesting
void sendTokenFileUpdatedMessage(InstanceType instanceType) {
    Criteria criteria = new Criteria();
    criteria.setInstanceName("%");
    criteria.setResource("%");
    criteria.setPartition("%");
    criteria.setPartitionState("%");
    criteria.setRecipientInstanceType(instanceType);
    /**
     * #HELIX-0.6.7-WORKAROUND
     * Add back when LIVESTANCES messaging is ported to 0.6 branch
     *    if (instanceType == InstanceType.PARTICIPANT) {
     *      criteria.setDataSource(Criteria.DataSource.LIVEINSTANCES);
     *    }
     */
    criteria.setSessionSpecific(true);
    Message tokenFileUpdatedMessage = new Message(Message.MessageType.USER_DEFINE_MSG, HelixMessageSubTypes.TOKEN_FILE_UPDATED.toString().toLowerCase() + UUID.randomUUID().toString());
    tokenFileUpdatedMessage.setMsgSubType(HelixMessageSubTypes.TOKEN_FILE_UPDATED.toString());
    tokenFileUpdatedMessage.setMsgState(Message.MessageState.NEW);
    if (instanceType == InstanceType.CONTROLLER) {
        tokenFileUpdatedMessage.setTgtSessionId("*");
    }
    // #HELIX-0.6.7-WORKAROUND
    // Temporarily bypass the default messaging service to allow upgrade to 0.6.7 which is missing support
    // for messaging to instances
    // int messagesSent = this.helixManager.getMessagingService().send(criteria, tokenFileUpdatedMessage);
    GobblinHelixMessagingService messagingService = new GobblinHelixMessagingService(this.helixManager);
    int messagesSent = messagingService.send(criteria, tokenFileUpdatedMessage);
    LOGGER.info(String.format("Sent %d token file updated message(s) to the %s", messagesSent, instanceType));
}
Also used : Message(org.apache.helix.model.Message) Criteria(org.apache.helix.Criteria) GobblinHelixMessagingService(org.apache.gobblin.cluster.GobblinHelixMessagingService) VisibleForTesting(com.google.common.annotations.VisibleForTesting)

Example 17 with Criteria

use of org.apache.helix.Criteria in project helix by apache.

the class TestSchedulerMessage method testSchedulerMsg.

@Test(dependsOnMethods = "testSchedulerZeroMsg")
public void testSchedulerMsg() throws Exception {
    _factory._results.clear();
    HelixManager manager = null;
    for (int i = 0; i < NODE_NR; i++) {
        _participants[i].getMessagingService().registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);
        // _startCMResultMap.get(hostDest)._manager;
        manager = _participants[i];
    }
    Message schedulerMessage = new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
    schedulerMessage.setTgtSessionId("*");
    schedulerMessage.setTgtName("CONTROLLER");
    // TODO: change it to "ADMIN" ?
    schedulerMessage.setSrcName("CONTROLLER");
    // schedulerMessage.getRecord().setSimpleField(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE,
    // "TestSchedulerMsg");
    // Template for the individual message sent to each participant
    Message msg = new Message(_factory.getMessageTypes().get(0), "Template");
    msg.setTgtSessionId("*");
    msg.setMsgState(MessageState.NEW);
    // Criteria to send individual messages
    Criteria cr = new Criteria();
    cr.setInstanceName("localhost_%");
    cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
    cr.setSessionSpecific(false);
    cr.setResource("%");
    cr.setPartition("%");
    ObjectMapper mapper = new ObjectMapper();
    SerializationConfig serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    StringWriter sw = new StringWriter();
    mapper.writeValue(sw, cr);
    String crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
    schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");
    HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
    Builder keyBuilder = helixDataAccessor.keyBuilder();
    helixDataAccessor.createControllerMessage(schedulerMessage);
    for (int i = 0; i < 30; i++) {
        Thread.sleep(2000);
        if (_PARTITIONS == _factory._results.size()) {
            break;
        }
    }
    Assert.assertEquals(_PARTITIONS, _factory._results.size());
    PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), schedulerMessage.getMsgId());
    int messageResultCount = 0;
    for (int i = 0; i < 10; i++) {
        Thread.sleep(1000);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount").equals("" + (_PARTITIONS * 3)));
        for (String key : statusUpdate.getMapFields().keySet()) {
            if (key.startsWith("MessageResult ")) {
                messageResultCount++;
                Assert.assertTrue(statusUpdate.getMapField(key).size() > 1);
            }
        }
        if (messageResultCount == _PARTITIONS * 3) {
            break;
        } else {
            Thread.sleep(2000);
        }
    }
    Assert.assertEquals(messageResultCount, _PARTITIONS * 3);
    int count = 0;
    for (Set<String> val : _factory._results.values()) {
        count += val.size();
    }
    Assert.assertEquals(count, _PARTITIONS * 3);
    // test the ZkPathDataDumpTask
    String controllerStatusPath = PropertyPathBuilder.controllerStatusUpdate(manager.getClusterName());
    List<String> subPaths = _gZkClient.getChildren(controllerStatusPath);
    Assert.assertTrue(subPaths.size() > 0);
    for (String subPath : subPaths) {
        String nextPath = controllerStatusPath + "/" + subPath;
        List<String> subsubPaths = _gZkClient.getChildren(nextPath);
        Assert.assertTrue(subsubPaths.size() > 0);
    }
    String instanceStatusPath = PropertyPathBuilder.instanceStatusUpdate(manager.getClusterName(), "localhost_" + (START_PORT));
    subPaths = _gZkClient.getChildren(instanceStatusPath);
    Assert.assertTrue(subPaths.size() == 0);
    for (String subPath : subPaths) {
        String nextPath = instanceStatusPath + "/" + subPath;
        List<String> subsubPaths = _gZkClient.getChildren(nextPath);
        Assert.assertTrue(subsubPaths.size() > 0);
        for (String subsubPath : subsubPaths) {
            String nextnextPath = nextPath + "/" + subsubPath;
            Assert.assertTrue(_gZkClient.getChildren(nextnextPath).size() > 0);
        }
    }
    Thread.sleep(3000);
    ZKPathDataDumpTask dumpTask = new ZKPathDataDumpTask(manager, 0L, 0L, Integer.MAX_VALUE);
    dumpTask.run();
    subPaths = _gZkClient.getChildren(controllerStatusPath);
    Assert.assertTrue(subPaths.size() > 0);
    for (String subPath : subPaths) {
        String nextPath = controllerStatusPath + "/" + subPath;
        List<String> subsubPaths = _gZkClient.getChildren(nextPath);
        Assert.assertTrue(subsubPaths.size() == 0);
    }
    subPaths = _gZkClient.getChildren(instanceStatusPath);
    Assert.assertTrue(subPaths.size() == 0);
    for (String subPath : subPaths) {
        String nextPath = instanceStatusPath + "/" + subPath;
        List<String> subsubPaths = _gZkClient.getChildren(nextPath);
        Assert.assertTrue(subsubPaths.size() > 0);
        for (String subsubPath : subsubPaths) {
            String nextnextPath = nextPath + "/" + subsubPath;
            Assert.assertTrue(_gZkClient.getChildren(nextnextPath).size() == 0);
        }
    }
}
Also used : HelixManager(org.apache.helix.HelixManager) Message(org.apache.helix.model.Message) SerializationConfig(org.codehaus.jackson.map.SerializationConfig) Builder(org.apache.helix.PropertyKey.Builder) PropertyPathBuilder(org.apache.helix.PropertyPathBuilder) Criteria(org.apache.helix.Criteria) ZKPathDataDumpTask(org.apache.helix.monitoring.ZKPathDataDumpTask) HelixDataAccessor(org.apache.helix.HelixDataAccessor) StringWriter(java.io.StringWriter) ObjectMapper(org.codehaus.jackson.map.ObjectMapper) PropertyKey(org.apache.helix.PropertyKey) ZNRecord(org.apache.helix.ZNRecord) Test(org.testng.annotations.Test)

Example 18 with Criteria

use of org.apache.helix.Criteria in project helix by apache.

the class TestSchedulerMessage method testSchedulerMsg3.

@Test(dependsOnMethods = "testSchedulerMsg")
public void testSchedulerMsg3() throws Exception {
    _factory._results.clear();
    Thread.sleep(2000);
    HelixManager manager = null;
    for (int i = 0; i < NODE_NR; i++) {
        _participants[i].getMessagingService().registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);
        _participants[i].getMessagingService().registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);
        // _startCMResultMap.get(hostDest)._manager;
        manager = _participants[i];
    }
    Message schedulerMessage = new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
    schedulerMessage.setTgtSessionId("*");
    schedulerMessage.setTgtName("CONTROLLER");
    // TODO: change it to "ADMIN" ?
    schedulerMessage.setSrcName("CONTROLLER");
    // Template for the individual message sent to each participant
    Message msg = new Message(_factory.getMessageTypes().get(0), "Template");
    msg.setTgtSessionId("*");
    msg.setMsgState(MessageState.NEW);
    // Criteria to send individual messages
    Criteria cr = new Criteria();
    cr.setInstanceName("localhost_%");
    cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
    cr.setSessionSpecific(false);
    cr.setResource("%");
    cr.setPartition("%");
    ObjectMapper mapper = new ObjectMapper();
    SerializationConfig serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    StringWriter sw = new StringWriter();
    mapper.writeValue(sw, cr);
    String crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
    schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");
    schedulerMessage.getRecord().setSimpleField("WAIT_ALL", "true");
    schedulerMessage.getRecord().setSimpleField(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsg3");
    Criteria cr2 = new Criteria();
    cr2.setRecipientInstanceType(InstanceType.CONTROLLER);
    cr2.setInstanceName("*");
    cr2.setSessionSpecific(false);
    MockAsyncCallback callback = new MockAsyncCallback();
    cr.setInstanceName("localhost_%");
    mapper = new ObjectMapper();
    serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    sw = new StringWriter();
    mapper.writeValue(sw, cr);
    crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    for (int i = 0; i < 4; i++) {
        callback = new MockAsyncCallback();
        cr.setInstanceName("localhost_" + (START_PORT + i));
        mapper = new ObjectMapper();
        serializationConfig = mapper.getSerializationConfig();
        serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
        sw = new StringWriter();
        mapper.writeValue(sw, cr);
        schedulerMessage.setMsgId(UUID.randomUUID().toString());
        crString = sw.toString();
        schedulerMessage.getRecord().setSimpleField("Criteria", crString);
        manager.getMessagingService().sendAndWait(cr2, schedulerMessage, callback, -1);
        String msgId = callback._message.getResultMap().get(DefaultSchedulerMessageHandlerFactory.SCHEDULER_MSG_ID);
        HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
        Builder keyBuilder = helixDataAccessor.keyBuilder();
        for (int j = 0; j < 100; j++) {
            Thread.sleep(200);
            PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
            ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
            if (statusUpdate.getMapFields().containsKey("Summary")) {
                break;
            }
        }
        Thread.sleep(3000);
        PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount").equals("" + (_PARTITIONS * 3 / 5)));
        int messageResultCount = 0;
        for (String key : statusUpdate.getMapFields().keySet()) {
            if (key.startsWith("MessageResult")) {
                messageResultCount++;
            }
        }
        Assert.assertEquals(messageResultCount, _PARTITIONS * 3 / 5);
        int count = 0;
        // System.out.println(i);
        for (Set<String> val : _factory._results.values()) {
            // System.out.println(val);
            count += val.size();
        }
        // System.out.println(count);
        Assert.assertEquals(count, _PARTITIONS * 3 / 5 * (i + 1));
    }
}
Also used : HelixManager(org.apache.helix.HelixManager) Message(org.apache.helix.model.Message) SerializationConfig(org.codehaus.jackson.map.SerializationConfig) Builder(org.apache.helix.PropertyKey.Builder) PropertyPathBuilder(org.apache.helix.PropertyPathBuilder) Criteria(org.apache.helix.Criteria) HelixDataAccessor(org.apache.helix.HelixDataAccessor) StringWriter(java.io.StringWriter) ObjectMapper(org.codehaus.jackson.map.ObjectMapper) PropertyKey(org.apache.helix.PropertyKey) ZNRecord(org.apache.helix.ZNRecord) Test(org.testng.annotations.Test)

Example 19 with Criteria

use of org.apache.helix.Criteria in project helix by apache.

the class TestSchedulerMessage method testSchedulerMsg4.

@Test(dependsOnMethods = "testSchedulerMsg3")
public void testSchedulerMsg4() throws Exception {
    _factory._results.clear();
    HelixManager manager = null;
    for (int i = 0; i < NODE_NR; i++) {
        _participants[i].getMessagingService().registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);
        _participants[i].getMessagingService().registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);
        // _startCMResultMap.get(hostDest)._manager;
        manager = _participants[i];
    }
    Message schedulerMessage = new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
    schedulerMessage.setTgtSessionId("*");
    schedulerMessage.setTgtName("CONTROLLER");
    // TODO: change it to "ADMIN" ?
    schedulerMessage.setSrcName("CONTROLLER");
    // Template for the individual message sent to each participant
    Message msg = new Message(_factory.getMessageTypes().get(0), "Template");
    msg.setTgtSessionId("*");
    msg.setMsgState(MessageState.NEW);
    // Criteria to send individual messages
    Criteria cr = new Criteria();
    cr.setInstanceName("localhost_%");
    cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
    cr.setSessionSpecific(false);
    cr.setResource("TestDB");
    cr.setPartition("%");
    ObjectMapper mapper = new ObjectMapper();
    SerializationConfig serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    StringWriter sw = new StringWriter();
    mapper.writeValue(sw, cr);
    String crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
    schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");
    schedulerMessage.getRecord().setSimpleField("WAIT_ALL", "true");
    schedulerMessage.getRecord().setSimpleField(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsg4");
    Criteria cr2 = new Criteria();
    cr2.setRecipientInstanceType(InstanceType.CONTROLLER);
    cr2.setInstanceName("*");
    cr2.setSessionSpecific(false);
    Map<String, String> constraints = new TreeMap<String, String>();
    constraints.put("MESSAGE_TYPE", "STATE_TRANSITION");
    constraints.put("TRANSITION", "OFFLINE-COMPLETED");
    constraints.put("CONSTRAINT_VALUE", "1");
    constraints.put("INSTANCE", ".*");
    manager.getClusterManagmentTool().setConstraint(manager.getClusterName(), ConstraintType.MESSAGE_CONSTRAINT, "constraint1", new ConstraintItem(constraints));
    MockAsyncCallback callback = new MockAsyncCallback();
    cr.setInstanceName("localhost_%");
    mapper = new ObjectMapper();
    serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    sw = new StringWriter();
    mapper.writeValue(sw, cr);
    crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    manager.getMessagingService().sendAndWait(cr2, schedulerMessage, callback, -1);
    String msgIdPrime = callback._message.getResultMap().get(DefaultSchedulerMessageHandlerFactory.SCHEDULER_MSG_ID);
    HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
    Builder keyBuilder = helixDataAccessor.keyBuilder();
    ArrayList<String> msgIds = new ArrayList<String>();
    for (int i = 0; i < NODE_NR; i++) {
        callback = new MockAsyncCallback();
        cr.setInstanceName("localhost_" + (START_PORT + i));
        mapper = new ObjectMapper();
        serializationConfig = mapper.getSerializationConfig();
        serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
        sw = new StringWriter();
        mapper.writeValue(sw, cr);
        schedulerMessage.setMsgId(UUID.randomUUID().toString());
        crString = sw.toString();
        schedulerMessage.getRecord().setSimpleField("Criteria", crString);
        manager.getMessagingService().sendAndWait(cr2, schedulerMessage, callback, -1);
        String msgId = callback._message.getResultMap().get(DefaultSchedulerMessageHandlerFactory.SCHEDULER_MSG_ID);
        msgIds.add(msgId);
    }
    for (int i = 0; i < NODE_NR; i++) {
        String msgId = msgIds.get(i);
        for (int j = 0; j < 100; j++) {
            Thread.sleep(200);
            PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
            ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
            if (statusUpdate.getMapFields().containsKey("Summary")) {
                // System.err.println(msgId+" done");
                break;
            }
        }
        PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount").equals("" + (_PARTITIONS * 3 / 5)));
        int messageResultCount = 0;
        for (String key : statusUpdate.getMapFields().keySet()) {
            if (key.startsWith("MessageResult")) {
                messageResultCount++;
            }
        }
        if (messageResultCount != _PARTITIONS * 3 / 5) {
            int x = 10;
            x = x + messageResultCount;
        }
        Assert.assertEquals(messageResultCount, _PARTITIONS * 3 / 5);
    }
    for (int j = 0; j < 100; j++) {
        Thread.sleep(200);
        PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgIdPrime);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        if (statusUpdate.getMapFields().containsKey("Summary")) {
            break;
        }
    }
    int count = 0;
    for (Set<String> val : _factory._results.values()) {
        // System.out.println(val);
        count += val.size();
    }
    // System.out.println(count);
    Assert.assertEquals(count, _PARTITIONS * 3 * 2);
}
Also used : HelixManager(org.apache.helix.HelixManager) Message(org.apache.helix.model.Message) SerializationConfig(org.codehaus.jackson.map.SerializationConfig) Builder(org.apache.helix.PropertyKey.Builder) PropertyPathBuilder(org.apache.helix.PropertyPathBuilder) ArrayList(java.util.ArrayList) Criteria(org.apache.helix.Criteria) TreeMap(java.util.TreeMap) HelixDataAccessor(org.apache.helix.HelixDataAccessor) StringWriter(java.io.StringWriter) ConstraintItem(org.apache.helix.model.ConstraintItem) ObjectMapper(org.codehaus.jackson.map.ObjectMapper) PropertyKey(org.apache.helix.PropertyKey) ZNRecord(org.apache.helix.ZNRecord) Test(org.testng.annotations.Test)

Example 20 with Criteria

use of org.apache.helix.Criteria in project helix by apache.

the class TestSchedulerMsgContraints method testSchedulerMsgContraints.

@Test
public void testSchedulerMsgContraints() throws Exception {
    TestSchedulerMessage.TestMessagingHandlerFactoryLatch factory = new TestSchedulerMessage.TestMessagingHandlerFactoryLatch();
    HelixManager manager = null;
    for (int i = 0; i < NODE_NR; i++) {
        _participants[i].getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
        _participants[i].getMessagingService().registerMessageHandlerFactory(factory.getMessageType(), factory);
        // _startCMResultMap.get(hostDest)._manager;
        manager = _participants[i];
    }
    Message schedulerMessage = new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
    schedulerMessage.setTgtSessionId("*");
    schedulerMessage.setTgtName("CONTROLLER");
    // TODO: change it to "ADMIN" ?
    schedulerMessage.setSrcName("CONTROLLER");
    // Template for the individual message sent to each participant
    Message msg = new Message(factory.getMessageType(), "Template");
    msg.setTgtSessionId("*");
    msg.setMsgState(MessageState.NEW);
    // Criteria to send individual messages
    Criteria cr = new Criteria();
    cr.setInstanceName("localhost_%");
    cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
    cr.setSessionSpecific(false);
    cr.setResource("%");
    cr.setPartition("%");
    ObjectMapper mapper = new ObjectMapper();
    SerializationConfig serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    StringWriter sw = new StringWriter();
    mapper.writeValue(sw, cr);
    String crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
    schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");
    schedulerMessage.getRecord().setSimpleField("WAIT_ALL", "true");
    schedulerMessage.getRecord().setSimpleField(DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsgContraints");
    Criteria cr2 = new Criteria();
    cr2.setRecipientInstanceType(InstanceType.CONTROLLER);
    cr2.setInstanceName("*");
    cr2.setSessionSpecific(false);
    TestSchedulerMessage.MockAsyncCallback callback = new TestSchedulerMessage.MockAsyncCallback();
    mapper = new ObjectMapper();
    serializationConfig = mapper.getSerializationConfig();
    serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
    sw = new StringWriter();
    mapper.writeValue(sw, cr);
    HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
    PropertyKey.Builder keyBuilder = helixDataAccessor.keyBuilder();
    // Set constraints that only 1 message per participant
    Map<String, String> constraints = new TreeMap<String, String>();
    constraints.put("MESSAGE_TYPE", "STATE_TRANSITION");
    constraints.put("TRANSITION", "OFFLINE-COMPLETED");
    constraints.put("CONSTRAINT_VALUE", "1");
    constraints.put("INSTANCE", ".*");
    manager.getClusterManagmentTool().setConstraint(manager.getClusterName(), ConstraintType.MESSAGE_CONSTRAINT, "constraint1", new ConstraintItem(constraints));
    // Send scheduler message
    crString = sw.toString();
    schedulerMessage.getRecord().setSimpleField("Criteria", crString);
    manager.getMessagingService().sendAndWait(cr2, schedulerMessage, callback, -1);
    String msgId = callback._message.getResultMap().get(DefaultSchedulerMessageHandlerFactory.SCHEDULER_MSG_ID);
    for (int j = 0; j < 10; j++) {
        Thread.sleep(200);
        PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        if (statusUpdate.getMapFields().containsKey("SentMessageCount")) {
            Assert.assertEquals(statusUpdate.getMapFields().get("SentMessageCount").get("MessageCount"), "" + (_PARTITIONS * 3));
            break;
        }
    }
    for (int i = 0; i < _PARTITIONS * 3 / 5; i++) {
        for (int j = 0; j < 10; j++) {
            Thread.sleep(300);
            if (factory._messageCount == 5 * (i + 1))
                break;
        }
        Thread.sleep(300);
        Assert.assertEquals(factory._messageCount, 5 * (i + 1));
        factory.signal();
    // System.err.println(i);
    }
    for (int j = 0; j < 10; j++) {
        Thread.sleep(200);
        PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
        ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
        if (statusUpdate.getMapFields().containsKey("Summary")) {
            break;
        }
    }
    Assert.assertEquals(_PARTITIONS, factory._results.size());
    PropertyKey controllerTaskStatus = keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), msgId);
    ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
    Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount").equals("" + (_PARTITIONS * 3)));
    int messageResultCount = 0;
    for (String key : statusUpdate.getMapFields().keySet()) {
        if (key.startsWith("MessageResult ")) {
            messageResultCount++;
        }
    }
    Assert.assertEquals(messageResultCount, _PARTITIONS * 3);
    int count = 0;
    for (Set<String> val : factory._results.values()) {
        count += val.size();
    }
    Assert.assertEquals(count, _PARTITIONS * 3);
    manager.getClusterManagmentTool().removeConstraint(manager.getClusterName(), ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
}
Also used : HelixManager(org.apache.helix.HelixManager) Message(org.apache.helix.model.Message) SerializationConfig(org.codehaus.jackson.map.SerializationConfig) Criteria(org.apache.helix.Criteria) TreeMap(java.util.TreeMap) HelixDataAccessor(org.apache.helix.HelixDataAccessor) StringWriter(java.io.StringWriter) ConstraintItem(org.apache.helix.model.ConstraintItem) ObjectMapper(org.codehaus.jackson.map.ObjectMapper) PropertyKey(org.apache.helix.PropertyKey) ZNRecord(org.apache.helix.ZNRecord) Test(org.testng.annotations.Test)

Aggregations

Criteria (org.apache.helix.Criteria)22 Message (org.apache.helix.model.Message)21 Test (org.testng.annotations.Test)15 HelixDataAccessor (org.apache.helix.HelixDataAccessor)8 HelixManager (org.apache.helix.HelixManager)8 PropertyKey (org.apache.helix.PropertyKey)8 StringWriter (java.io.StringWriter)7 ZNRecord (org.apache.helix.ZNRecord)7 ObjectMapper (org.codehaus.jackson.map.ObjectMapper)7 SerializationConfig (org.codehaus.jackson.map.SerializationConfig)7 UUID (java.util.UUID)6 Builder (org.apache.helix.PropertyKey.Builder)6 VisibleForTesting (com.google.common.annotations.VisibleForTesting)5 PropertyPathBuilder (org.apache.helix.PropertyPathBuilder)4 AsyncCallback (org.apache.helix.messaging.AsyncCallback)4 TreeMap (java.util.TreeMap)2 ConstraintItem (org.apache.helix.model.ConstraintItem)2 SegmentRefreshMessage (com.linkedin.pinot.common.messages.SegmentRefreshMessage)1 ArrayList (java.util.ArrayList)1 GobblinHelixMessagingService (org.apache.gobblin.cluster.GobblinHelixMessagingService)1