<commit_before>package org.apache.helix.actor;

import org.apache.helix.*;
import org.apache.helix.controller.HelixControllerMain;
import org.apache.helix.model.HelixConfigScope;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.Message;
import org.apache.helix.model.Partition;
import org.apache.helix.participant.statemachine.StateModel;
import org.apache.helix.participant.statemachine.StateModelFactory;
import org.apache.helix.participant.statemachine.StateModelInfo;
import org.apache.helix.participant.statemachine.Transition;
import org.apache.helix.tools.ClusterSetup;
import org.apache.helix.tools.ClusterStateVerifier;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

public class TestNettyHelixActor extends ZkUnitTestBase {

    private static final String CLUSTER_NAME = "TEST_CLUSTER";
    private static final String RESOURCE_NAME = "MyResource";
    private static final HelixActorMessageCodec<String> CODEC = new HelixActorMessageCodec<String>() {
        @Override
        public byte[] encode(String message) {
            return message.getBytes();
        }

        @Override
        public String decode(byte[] message) {
            return new String(message);
        }
    };

    private int firstPort;
    private int secondPort;
    private HelixManager controller;
    private HelixManager firstNode;
    private HelixManager secondNode;

    @BeforeClass
    public void beforeClass() throws Exception {
        // Allocate test resources
        firstPort = TestHelper.getRandomPort();
        secondPort = TestHelper.getRandomPort();

        // Setup cluster
        ClusterSetup clusterSetup = new ClusterSetup(ZK_ADDR);
        clusterSetup.addCluster(CLUSTER_NAME, true);
        clusterSetup.addInstanceToCluster(CLUSTER_NAME, "localhost_" + firstPort);
        clusterSetup.addInstanceToCluster(CLUSTER_NAME, "localhost_" + secondPort);
        clusterSetup.setConfig(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, CLUSTER_NAME + ",localhost_" + firstPort, "ACTOR_PORT=" + firstPort);
        clusterSetup.setConfig(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, CLUSTER_NAME + ",localhost_" + secondPort, "ACTOR_PORT=" + secondPort);

        // Start Helix agents
        controller = HelixControllerMain.startHelixController(ZK_ADDR, CLUSTER_NAME, "CONTROLLER", "STANDALONE");
        firstNode = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "localhost_" + firstPort, InstanceType.PARTICIPANT, ZK_ADDR);
        secondNode = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "localhost_" + secondPort, InstanceType.PARTICIPANT, ZK_ADDR);

        // Connect participants
        firstNode.getStateMachineEngine().registerStateModelFactory("OnlineOffline", new DummyStateModelFactory());
        firstNode.connect();
        secondNode.getStateMachineEngine().registerStateModelFactory("OnlineOffline", new DummyStateModelFactory());
        secondNode.connect();

        // Add a resource
        clusterSetup.addResourceToCluster(CLUSTER_NAME, RESOURCE_NAME, 4, "OnlineOffline");
        clusterSetup.rebalanceResource(CLUSTER_NAME, RESOURCE_NAME, 1);

        // Wait for External view convergence
        ClusterStateVerifier.verifyByZkCallback(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME), 10000);
    }

    @AfterClass
    public void afterClass() throws Exception {
        firstNode.disconnect();
        secondNode.disconnect();
        controller.disconnect();
    }

    @Test
    public void testMessagePassing() throws Exception {
        // Start first Actor w/ counter
        final ConcurrentMap<String, AtomicInteger> firstCounts = new ConcurrentHashMap<String, AtomicInteger>();
        NettyHelixActor<String> firstActor = new NettyHelixActor<String>(firstNode, firstPort, CODEC);
        firstActor.register(RESOURCE_NAME, new HelixActorCallback<String>() {
            @Override
            public void onMessage(Partition partition, String state, String message) {
                String key = partition.getPartitionName() + ":" + state;
                firstCounts.putIfAbsent(key, new AtomicInteger());
                firstCounts.get(key).incrementAndGet();
            }
        });
        firstNode.addExternalViewChangeListener(firstActor);
        firstActor.start();

        // Start second Actor w/ counter
        final ConcurrentMap<String, AtomicInteger> secondCounts = new ConcurrentHashMap<String, AtomicInteger>();
        NettyHelixActor<String> secondActor = new NettyHelixActor<String>(secondNode, secondPort, CODEC);
        secondActor.register(RESOURCE_NAME, new HelixActorCallback<String>() {
            @Override
            public void onMessage(Partition partition, String state, String message) {
                String key = partition.getPartitionName() + ":" + state;
                secondCounts.putIfAbsent(key, new AtomicInteger());
                secondCounts.get(key).incrementAndGet();
            }
        });
        secondNode.addExternalViewChangeListener(secondActor);
        secondActor.start();

        // Find all partitions on second node...
        String secondName = "localhost_" + secondPort;
        Set<String> secondPartitions = new HashSet<String>();
        IdealState idealState = controller.getClusterManagmentTool().getResourceIdealState(CLUSTER_NAME, RESOURCE_NAME);
        for (String partitionName : idealState.getPartitionSet()) {
            for (Map.Entry<String, String> stateEntry : idealState.getInstanceStateMap(partitionName).entrySet()) {
                if (stateEntry.getKey().equals(secondName)) {
                    secondPartitions.add(partitionName);
                }
            }
        }

        // And use first node to send messages to them
        for (String partitionName : secondPartitions) {
            firstActor.send(new Partition(partitionName), "ONLINE", "Hello world!");
        }

        // Loopback
        for (String partitionName : secondPartitions) {
            secondActor.send(new Partition(partitionName), "ONLINE", "Hello world!");
        }

        // Check
        Thread.sleep(100); // just in case
        for (String partitionName : secondPartitions) {
            AtomicInteger count = secondCounts.get(partitionName + ":ONLINE");
            Assert.assertNotNull(count);
            Assert.assertEquals(count.get(), 2);
        }
    }

    public static class DummyStateModelFactory extends StateModelFactory<StateModel> {
        @Override
        public StateModel createNewStateModel(String partitionName) {<commit_msg>send some more messages in test just because<commit_after>package org.apache.helix.actor;

import org.apache.helix.*;
import org.apache.helix.controller.HelixControllerMain;
import org.apache.helix.model.HelixConfigScope;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.Message;
import org.apache.helix.model.Partition;
import org.apache.helix.participant.statemachine.StateModel;
import org.apache.helix.participant.statemachine.StateModelFactory;
import org.apache.helix.participant.statemachine.StateModelInfo;
import org.apache.helix.participant.statemachine.Transition;
import org.apache.helix.tools.ClusterSetup;
import org.apache.helix.tools.ClusterStateVerifier;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

public class TestNettyHelixActor extends ZkUnitTestBase {

    private static final String CLUSTER_NAME = "TEST_CLUSTER";
    private static final String RESOURCE_NAME = "MyResource";
    private static final HelixActorMessageCodec<String> CODEC = new HelixActorMessageCodec<String>() {
        @Override
        public byte[] encode(String message) {
            return message.getBytes();
        }

        @Override
        public String decode(byte[] message) {
            return new String(message);
        }
    };

    private int firstPort;
    private int secondPort;
    private HelixManager controller;
    private HelixManager firstNode;
    private HelixManager secondNode;

    @BeforeClass
    public void beforeClass() throws Exception {
        // Allocate test resources
        firstPort = TestHelper.getRandomPort();
        secondPort = TestHelper.getRandomPort();

        // Setup cluster
        ClusterSetup clusterSetup = new ClusterSetup(ZK_ADDR);
        clusterSetup.addCluster(CLUSTER_NAME, true);
        clusterSetup.addInstanceToCluster(CLUSTER_NAME, "localhost_" + firstPort);
        clusterSetup.addInstanceToCluster(CLUSTER_NAME, "localhost_" + secondPort);
        clusterSetup.setConfig(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, CLUSTER_NAME + ",localhost_" + firstPort, "ACTOR_PORT=" + firstPort);
        clusterSetup.setConfig(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, CLUSTER_NAME + ",localhost_" + secondPort, "ACTOR_PORT=" + secondPort);

        // Start Helix agents
        controller = HelixControllerMain.startHelixController(ZK_ADDR, CLUSTER_NAME, "CONTROLLER", "STANDALONE");
        firstNode = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "localhost_" + firstPort, InstanceType.PARTICIPANT, ZK_ADDR);
        secondNode = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "localhost_" + secondPort, InstanceType.PARTICIPANT, ZK_ADDR);

        // Connect participants
        firstNode.getStateMachineEngine().registerStateModelFactory("OnlineOffline", new DummyStateModelFactory());
        firstNode.connect();
        secondNode.getStateMachineEngine().registerStateModelFactory("OnlineOffline", new DummyStateModelFactory());
        secondNode.connect();

        // Add a resource
        clusterSetup.addResourceToCluster(CLUSTER_NAME, RESOURCE_NAME, 4, "OnlineOffline");
        clusterSetup.rebalanceResource(CLUSTER_NAME, RESOURCE_NAME, 1);

        // Wait for External view convergence
        ClusterStateVerifier.verifyByZkCallback(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME), 10000);
    }

    @AfterClass
    public void afterClass() throws Exception {
        firstNode.disconnect();
        secondNode.disconnect();
        controller.disconnect();
    }

    @Test
    public void testMessagePassing() throws Exception {
        int numMessages = 100;

        // Start first Actor w/ counter
        final ConcurrentMap<String, AtomicInteger> firstCounts = new ConcurrentHashMap<String, AtomicInteger>();
        NettyHelixActor<String> firstActor = new NettyHelixActor<String>(firstNode, firstPort, CODEC);
        firstActor.register(RESOURCE_NAME, new HelixActorCallback<String>() {
            @Override
            public void onMessage(Partition partition, String state, String message) {
                String key = partition.getPartitionName() + ":" + state;
                firstCounts.putIfAbsent(key, new AtomicInteger());
                firstCounts.get(key).incrementAndGet();
            }
        });
        firstNode.addExternalViewChangeListener(firstActor);
        firstActor.start();

        // Start second Actor w/ counter
        final ConcurrentMap<String, AtomicInteger> secondCounts = new ConcurrentHashMap<String, AtomicInteger>();
        NettyHelixActor<String> secondActor = new NettyHelixActor<String>(secondNode, secondPort, CODEC);
        secondActor.register(RESOURCE_NAME, new HelixActorCallback<String>() {
            @Override
            public void onMessage(Partition partition, String state, String message) {
                String key = partition.getPartitionName() + ":" + state;
                secondCounts.putIfAbsent(key, new AtomicInteger());
                secondCounts.get(key).incrementAndGet();
            }
        });
        secondNode.addExternalViewChangeListener(secondActor);
        secondActor.start();

        // Find all partitions on second node...
        String secondName = "localhost_" + secondPort;
        Set<String> secondPartitions = new HashSet<String>();
        IdealState idealState = controller.getClusterManagmentTool().getResourceIdealState(CLUSTER_NAME, RESOURCE_NAME);
        for (String partitionName : idealState.getPartitionSet()) {
            for (Map.Entry<String, String> stateEntry : idealState.getInstanceStateMap(partitionName).entrySet()) {
                if (stateEntry.getKey().equals(secondName)) {
                    secondPartitions.add(partitionName);
                }
            }
        }

        // And use first node to send messages to them
        for (String partitionName : secondPartitions) {
            for (int i = 0; i < numMessages; i++) {
                firstActor.send(new Partition(partitionName), "ONLINE", "Hello world!");
            }
        }

        // Loopback
        for (String partitionName : secondPartitions) {
            for (int i = 0; i < numMessages; i++) {
                secondActor.send(new Partition(partitionName), "ONLINE", "Hello world!");
            }
        }

        // Check
        Thread.sleep(100); // just in case
        for (String partitionName : secondPartitions) {
            AtomicInteger count = secondCounts.get(partitionName + ":ONLINE");
            Assert.assertNotNull(count);
            Assert.assertEquals(count.get(), 2 * numMessages);
        }
    }

    public static class DummyStateModelFactory extends StateModelFactory<StateModel> {
        @Override
        public StateModel createNewStateModel(String partitionName) {
{'fraction_non_alphanumeric': 0.09367631296891747, 'fraction_numerical': 0.0019292604501607716, 'mean_word_length': 3.310440406529104, 'pattern_counts': {'":': 6, '<': 37, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1794544', 'n_tokens_mistral': 3813, 'n_tokens_neox': 3674, 'n_words': 761}