<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.pulsar.broker.loadbalance;

import static org.apache.pulsar.broker.cache.ConfigurationCacheService.POLICIES;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertTrue;

import com.google.common.collect.Lists;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.bookkeeper.util.ZkUtils;
import org.apache.pulsar.broker.PulsarService;
import org.apache.pulsar.broker.ServiceConfiguration;
import org.apache.pulsar.broker.admin.AdminResource;
import org.apache.pulsar.broker.loadbalance.impl.PulsarResourceDescription;
import org.apache.pulsar.broker.loadbalance.impl.SimpleLoadManagerImpl;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.internal.NamespacesImpl;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyData;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyDataImpl;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyType;
import org.apache.pulsar.common.policies.data.BundlesData;
import org.apache.pulsar.common.policies.data.NamespaceIsolationData;
import org.apache.pulsar.common.policies.data.NamespaceIsolationDataImpl;
import org.apache.pulsar.common.policies.data.Policies;
import org.apache.pulsar.common.policies.data.ResourceQuota;
import org.apache.pulsar.common.policies.impl.NamespaceIsolationPolicies;
import org.apache.pulsar.common.util.ObjectMapperFactory;
import org.apache.pulsar.metadata.api.MetadataStoreException.BadVersionException;
import org.apache.pulsar.metadata.api.coordination.ResourceLock;
import org.apache.pulsar.policies.data.loadbalancer.LoadReport;
import org.apache.pulsar.policies.data.loadbalancer.NamespaceBundleStats;
import org.apache.pulsar.policies.data.loadbalancer.ResourceUsage;
import org.apache.pulsar.policies.data.loadbalancer.SystemResourceUsage;
import org.apache.pulsar.zookeeper.LocalBookkeeperEnsemble;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooDefs;
import org.awaitility.Awaitility;
import org.powermock.reflect.Whitebox;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

/**
 * Start two brokers in the same cluster and have them connect to the same zookeeper. When the PulsarService starts, it
 * will do the leader election and one of the brokers will become the leader. Then kill that broker and check if the
 * second one becomes the leader.
 */
@Test(groups = "broker")
public class LoadBalancerTest {
    LocalBookkeeperEnsemble bkEnsemble;

    ExecutorService executor = new ThreadPoolExecutor(5, 20, 30, TimeUnit.SECONDS, new LinkedBlockingQueue<>());

    private static final Logger log = LoggerFactory.getLogger(LoadBalancerTest.class);

    private static final int MAX_RETRIES = 15;

    private static final int BROKER_COUNT = 5;
    private int[] brokerWebServicePorts = new int[BROKER_COUNT];
    private int[] brokerNativeBrokerPorts = new int[BROKER_COUNT];
    private URL[] brokerUrls = new URL[BROKER_COUNT];
    private String[] lookupAddresses = new String[BROKER_COUNT];
    private PulsarService[] pulsarServices = new PulsarService[BROKER_COUNT];
    private PulsarAdmin[] pulsarAdmins = new PulsarAdmin[BROKER_COUNT];

    @BeforeMethod
    void setup() throws Exception {
        // Start local bookkeeper ensemble
        bkEnsemble = new LocalBookkeeperEnsemble(3, 0, () -> 0);
        bkEnsemble.start();
        ZkUtils.createFullPathOptimistic(bkEnsemble.getZkClient(),
                SimpleLoadManagerImpl.LOADBALANCER_DYNAMIC_SETTING_STRATEGY_ZPATH,
                "{\"loadBalancerStrategy\":\"leastLoadedServer\"}".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT);

        final String localhost = "localhost";
        // start brokers
        for (int i = 0; i < BROKER_COUNT; i++) {


            ServiceConfiguration config = new ServiceConfiguration();
            config.setBrokerServicePort(Optional.of(brokerNativeBrokerPorts[i]));
            config.setClusterName("use");
            config.setAdvertisedAddress(localhost);
            config.setAdvertisedAddress("localhost");
            config.setWebServicePort(Optional.of(0));
            config.setBrokerServicePortTls(Optional.of(0));
            config.setWebServicePortTls(Optional.of(0));
            config.setZookeeperServers("127.0.0.1" + ":" + bkEnsemble.getZookeeperPort());
            config.setBrokerShutdownTimeoutMs(0L);
            config.setBrokerServicePort(Optional.of(0));
            config.setLoadManagerClassName(SimpleLoadManagerImpl.class.getName());
            config.setAdvertisedAddress(localhost+i);
            config.setLoadBalancerEnabled(false);

            pulsarServices[i] = new PulsarService(config);
            pulsarServices[i].start();
            brokerWebServicePorts[i] = pulsarServices[i].getListenPortHTTP().get();
            brokerNativeBrokerPorts[i] = pulsarServices[i].getBrokerListenPort().get();

            brokerUrls[i] = new URL("http://127.0.0.1" + ":" + brokerWebServicePorts[i]);
            lookupAddresses[i] = pulsarServices[i].getAdvertisedAddress() + ":" + pulsarServices[i].getListenPortHTTP().get();
            pulsarAdmins[i] = PulsarAdmin.builder().serviceHttpUrl(brokerUrls[i].toString()).build();
        }

        createNamespacePolicies(pulsarServices[0]);

        Thread.sleep(100);
    }

    @AfterMethod(alwaysRun = true)
    void shutdown() throws Exception {
        log.info("--- Shutting down ---");
        executor.shutdownNow();

        for (int i = 0; i < BROKER_COUNT; i++) {
            pulsarAdmins[i].close();
            if (pulsarServices[i] != null) {
                pulsarServices[i].close();
            }
        }

        bkEnsemble.stop();
    }

    private void loopUntilLeaderChangesForAllBroker(List<PulsarService> activePulsars, LeaderBroker oldLeader) {
        int loopCount = 0;
        Awaitility.await()
            .pollInterval(1, TimeUnit.SECONDS)
            .atMost(MAX_RETRIES, TimeUnit.SECONDS)
            .until(() -> {
                boolean settled = true;
                // Check if the all active pulsar see a new leader.
                for (PulsarService pulsar : activePulsars) {
                    Optional<LeaderBroker> leader = pulsar.getLeaderElectionService().readCurrentLeader().join();
                    // Check leader a pulsar see is not empty and not the old leader.
                    if (!leader.isPresent() || leader.get().equals(oldLeader)) {
                        settled = false;
                        break;
                    }
                }
                return settled;
            });
        // Check if maximum retries are already done. If yes, assert.
        Assert.assertNotEquals(loopCount, MAX_RETRIES, "Leader is not changed even after maximum retries.");
    }

    /*
     * tests that load manager creates its node and writes the initial load report a /loadbalance/brokers tests that
     * those load reports can be deserialized and are in valid format tests if the rankings are populated from the load
     * reports are not, both broker will have zero rank
     */
    @SuppressWarnings("unchecked")
    @Test
    public void testLoadReportsWrittenOnMetadataStore() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            String path = String.format("%s/%s", SimpleLoadManagerImpl.LOADBALANCE_BROKERS_ROOT,
                    lookupAddresses[i]);
            byte[] loadReportData = pulsarServices[i].getLocalMetadataStore().get(path).join().get().getValue();
            assertTrue(loadReportData.length > 0);
            log.info("LoadReport {}, {}", lookupAddresses[i], new String(loadReportData));

            LoadReport loadReport = ObjectMapperFactory.getThreadLocal().readValue(loadReportData, LoadReport.class);
            assertEquals(loadReport.getName(), lookupAddresses[i]);

            // Check Initial Ranking is populated in both the brokers
            Field ranking = ((SimpleLoadManagerImpl) pulsarServices[i].getLoadManager().get()).getClass()
                    .getDeclaredField("sortedRankings");
            ranking.setAccessible(true);
            AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking = (AtomicReference<Map<Long, Set<ResourceUnit>>>) ranking
                    .get(pulsarServices[i].getLoadManager().get());
            printSortedRanking(sortedRanking);

            // all brokers have same rank to it would be 0 --> set-of-all-the-brokers
            int brokerCount = 0;
            for (Map.Entry<Long, Set<ResourceUnit>> entry : sortedRanking.get().entrySet()) {
                brokerCount += entry.getValue().size();
            }
            assertEquals(brokerCount, BROKER_COUNT);
            TopicName topicName = TopicName.get("persistent://pulsar/use/primary-ns/test-topic");
            ResourceUnit found = pulsarServices[i].getLoadManager().get()
                    .getLeastLoaded(pulsarServices[i].getNamespaceService().getBundle(topicName)).get();
            assertNotNull(found);
        }
    }

    /*
     * tests rankings get updated when we write write the new load reports to the zookeeper on loadbalance root node
     * tests writing pre-configured load report on the zookeeper translates the pre-calculated rankings
     */
    @Test
    public void testUpdateLoadReportAndCheckUpdatedRanking() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            LoadReport lr = new LoadReport();
            lr.setName(lookupAddresses[i]);
            SystemResourceUsage sru = new SystemResourceUsage();
            sru.setBandwidthIn(new ResourceUsage(256, 1024000));
            sru.setBandwidthOut(new ResourceUsage(250, 1024000));
            sru.setMemory(new ResourceUsage(1024, 8192));
            sru.setCpu(new ResourceUsage(5, 400));
            lr.setSystemResourceUsage(sru);

            Whitebox.setInternalState(pulsarServices[0].getLoadManager().get(), "lastLoadReport", lr);
            ResourceLock<LoadReport> lock = Whitebox.getInternalState(pulsarServices[i].getLoadManager().get(),
                    "brokerLock");
            lock.updateValue(lr).join();
        }

        for (int i = 0; i < BROKER_COUNT; i++) {
            Method updateRanking = Whitebox.getMethod(SimpleLoadManagerImpl.class, "updateRanking");
            updateRanking.invoke(pulsarServices[0].getLoadManager().get());
        }

        // do lookup for bunch of bundles
        int totalNamespaces = 200;
        Map<String, Integer> namespaceOwner = new HashMap<>();
        for (int i = 0; i < totalNamespaces; i++) {
            TopicName topicName = TopicName.get("persistent://pulsar/use/primary-ns-" + i + "/test-topic");
            ResourceUnit found = pulsarServices[0].getLoadManager().get()
                    .getLeastLoaded(pulsarServices[0].getNamespaceService().getBundle(topicName)).get();
            if (namespaceOwner.containsKey(found.getResourceId())) {
                namespaceOwner.put(found.getResourceId(), namespaceOwner.get(found.getResourceId()) + 1);
            } else {
                namespaceOwner.put(found.getResourceId(), 1);
            }
        }
        // assert that distribution variation is not more than 20%
        int averageNamespaces = totalNamespaces / BROKER_COUNT;
        int tenPercentOfAverageNamespaces = averageNamespaces / 10;
        int lowerBound = averageNamespaces - tenPercentOfAverageNamespaces;
        int upperBound = averageNamespaces + tenPercentOfAverageNamespaces;

        // assert each broker received ownership of fair amount of namespaces 90%+
        for (Map.Entry<String, Integer> broker : namespaceOwner.entrySet()) {
            log.info("Count of bundles assigned: {}, {} -- lower-bound: {} - upper-bound: {} ",
                    broker.getKey(), broker.getValue(), lowerBound, upperBound);
            assertTrue(broker.getValue() >= lowerBound && broker.getValue() <= upperBound);
        }
    }

    private AtomicReference<Map<Long, Set<ResourceUnit>>> getSortedRanking(PulsarService pulsar)
            throws NoSuchFieldException, IllegalAccessException {
        Field ranking = ((SimpleLoadManagerImpl) pulsar.getLoadManager().get()).getClass()
                .getDeclaredField("sortedRankings");
        ranking.setAccessible(true);
        @SuppressWarnings("unchecked")
        AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking = (AtomicReference<Map<Long, Set<ResourceUnit>>>) ranking
                .get(pulsar.getLoadManager().get());
        return sortedRanking;
    }

    private void printSortedRanking(AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking) {
        log.info("Sorted Ranking Result:");
        sortedRanking.get().forEach((score, rus) -> {
            for (ResourceUnit ru : rus) {
                log.info("  - {}, {}", ru.getResourceId(), score);
            }
        });
    }

    /*
     * Pre-publish load report to ZK, each broker has: - Difference memory capacity, for the first 3 brokers memory is
     * bottleneck, for the 4/5th brokers CPU become bottleneck since memory is big enough - non-bundles assigned so all
     * idle resources are available for new bundle Check the broker rankings are the load percentage of each broker.
     */
    @Test
    public void testBrokerRanking() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            LoadReport lr = new LoadReport();
            lr.setName(lookupAddresses[i]);
            SystemResourceUsage sru = new SystemResourceUsage();
            sru.setBandwidthIn(new ResourceUsage(0, 1024000));
            sru.setBandwidthOut(new ResourceUsage(0, 1024000));
            sru.setMemory(new ResourceUsage(1024, 2048 * (i + 1)));
            sru.setCpu(new ResourceUsage(60, 400));
            lr.setSystemResourceUsage(sru);

            ResourceLock<LoadReport> lock = Whitebox.getInternalState(pulsarServices[i].getLoadManager().get(),
                    "brokerLock");
            lock.updateValue(lr).join();
        }

        for (int i = 0; i < BROKER_COUNT; i++) {
            Method updateRanking = Whitebox.getMethod(SimpleLoadManagerImpl.class, "updateRanking");
            updateRanking.invoke(pulsarServices[0].getLoadManager().get());
        }

        // check the ranking result<commit_msg>fix flaky test testBrokerRanking (#11114)<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.pulsar.broker.loadbalance;

import static org.apache.pulsar.broker.cache.ConfigurationCacheService.POLICIES;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertTrue;

import com.google.common.collect.Lists;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.bookkeeper.util.ZkUtils;
import org.apache.pulsar.broker.PulsarService;
import org.apache.pulsar.broker.ServiceConfiguration;
import org.apache.pulsar.broker.admin.AdminResource;
import org.apache.pulsar.broker.loadbalance.impl.PulsarResourceDescription;
import org.apache.pulsar.broker.loadbalance.impl.SimpleLoadManagerImpl;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.internal.NamespacesImpl;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyData;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyDataImpl;
import org.apache.pulsar.common.policies.data.AutoFailoverPolicyType;
import org.apache.pulsar.common.policies.data.BundlesData;
import org.apache.pulsar.common.policies.data.NamespaceIsolationData;
import org.apache.pulsar.common.policies.data.NamespaceIsolationDataImpl;
import org.apache.pulsar.common.policies.data.Policies;
import org.apache.pulsar.common.policies.data.ResourceQuota;
import org.apache.pulsar.common.policies.impl.NamespaceIsolationPolicies;
import org.apache.pulsar.common.util.ObjectMapperFactory;
import org.apache.pulsar.metadata.api.MetadataStoreException.BadVersionException;
import org.apache.pulsar.metadata.api.coordination.ResourceLock;
import org.apache.pulsar.policies.data.loadbalancer.LoadReport;
import org.apache.pulsar.policies.data.loadbalancer.NamespaceBundleStats;
import org.apache.pulsar.policies.data.loadbalancer.ResourceUsage;
import org.apache.pulsar.policies.data.loadbalancer.SystemResourceUsage;
import org.apache.pulsar.zookeeper.LocalBookkeeperEnsemble;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooDefs;
import org.awaitility.Awaitility;
import org.powermock.reflect.Whitebox;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

/**
 * Start two brokers in the same cluster and have them connect to the same zookeeper. When the PulsarService starts, it
 * will do the leader election and one of the brokers will become the leader. Then kill that broker and check if the
 * second one becomes the leader.
 */
@Test(groups = "broker")
public class LoadBalancerTest {
    LocalBookkeeperEnsemble bkEnsemble;

    ExecutorService executor = new ThreadPoolExecutor(5, 20, 30, TimeUnit.SECONDS, new LinkedBlockingQueue<>());

    private static final Logger log = LoggerFactory.getLogger(LoadBalancerTest.class);

    private static final int MAX_RETRIES = 15;

    private static final int BROKER_COUNT = 5;
    private int[] brokerWebServicePorts = new int[BROKER_COUNT];
    private int[] brokerNativeBrokerPorts = new int[BROKER_COUNT];
    private URL[] brokerUrls = new URL[BROKER_COUNT];
    private String[] lookupAddresses = new String[BROKER_COUNT];
    private PulsarService[] pulsarServices = new PulsarService[BROKER_COUNT];
    private PulsarAdmin[] pulsarAdmins = new PulsarAdmin[BROKER_COUNT];

    @BeforeMethod
    void setup() throws Exception {
        // Start local bookkeeper ensemble
        bkEnsemble = new LocalBookkeeperEnsemble(3, 0, () -> 0);
        bkEnsemble.start();
        ZkUtils.createFullPathOptimistic(bkEnsemble.getZkClient(),
                SimpleLoadManagerImpl.LOADBALANCER_DYNAMIC_SETTING_STRATEGY_ZPATH,
                "{\"loadBalancerStrategy\":\"leastLoadedServer\"}".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT);

        final String localhost = "localhost";
        // start brokers
        for (int i = 0; i < BROKER_COUNT; i++) {


            ServiceConfiguration config = new ServiceConfiguration();
            config.setBrokerServicePort(Optional.of(brokerNativeBrokerPorts[i]));
            config.setClusterName("use");
            config.setAdvertisedAddress(localhost);
            config.setAdvertisedAddress("localhost");
            config.setWebServicePort(Optional.of(0));
            config.setBrokerServicePortTls(Optional.of(0));
            config.setWebServicePortTls(Optional.of(0));
            config.setZookeeperServers("127.0.0.1" + ":" + bkEnsemble.getZookeeperPort());
            config.setBrokerShutdownTimeoutMs(0L);
            config.setBrokerServicePort(Optional.of(0));
            config.setLoadManagerClassName(SimpleLoadManagerImpl.class.getName());
            config.setAdvertisedAddress(localhost+i);
            config.setLoadBalancerEnabled(false);

            pulsarServices[i] = new PulsarService(config);
            pulsarServices[i].start();
            brokerWebServicePorts[i] = pulsarServices[i].getListenPortHTTP().get();
            brokerNativeBrokerPorts[i] = pulsarServices[i].getBrokerListenPort().get();

            brokerUrls[i] = new URL("http://127.0.0.1" + ":" + brokerWebServicePorts[i]);
            lookupAddresses[i] = pulsarServices[i].getAdvertisedAddress() + ":" + pulsarServices[i].getListenPortHTTP().get();
            pulsarAdmins[i] = PulsarAdmin.builder().serviceHttpUrl(brokerUrls[i].toString()).build();
        }

        createNamespacePolicies(pulsarServices[0]);

        Thread.sleep(100);
    }

    @AfterMethod(alwaysRun = true)
    void shutdown() throws Exception {
        log.info("--- Shutting down ---");
        executor.shutdownNow();

        for (int i = 0; i < BROKER_COUNT; i++) {
            pulsarAdmins[i].close();
            if (pulsarServices[i] != null) {
                pulsarServices[i].close();
            }
        }

        bkEnsemble.stop();
    }

    private void loopUntilLeaderChangesForAllBroker(List<PulsarService> activePulsars, LeaderBroker oldLeader) {
        int loopCount = 0;
        Awaitility.await()
            .pollInterval(1, TimeUnit.SECONDS)
            .atMost(MAX_RETRIES, TimeUnit.SECONDS)
            .until(() -> {
                boolean settled = true;
                // Check if the all active pulsar see a new leader.
                for (PulsarService pulsar : activePulsars) {
                    Optional<LeaderBroker> leader = pulsar.getLeaderElectionService().readCurrentLeader().join();
                    // Check leader a pulsar see is not empty and not the old leader.
                    if (!leader.isPresent() || leader.get().equals(oldLeader)) {
                        settled = false;
                        break;
                    }
                }
                return settled;
            });
        // Check if maximum retries are already done. If yes, assert.
        Assert.assertNotEquals(loopCount, MAX_RETRIES, "Leader is not changed even after maximum retries.");
    }

    /*
     * tests that load manager creates its node and writes the initial load report a /loadbalance/brokers tests that
     * those load reports can be deserialized and are in valid format tests if the rankings are populated from the load
     * reports are not, both broker will have zero rank
     */
    @SuppressWarnings("unchecked")
    @Test
    public void testLoadReportsWrittenOnMetadataStore() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            String path = String.format("%s/%s", SimpleLoadManagerImpl.LOADBALANCE_BROKERS_ROOT,
                    lookupAddresses[i]);
            byte[] loadReportData = pulsarServices[i].getLocalMetadataStore().get(path).join().get().getValue();
            assertTrue(loadReportData.length > 0);
            log.info("LoadReport {}, {}", lookupAddresses[i], new String(loadReportData));

            LoadReport loadReport = ObjectMapperFactory.getThreadLocal().readValue(loadReportData, LoadReport.class);
            assertEquals(loadReport.getName(), lookupAddresses[i]);

            // Check Initial Ranking is populated in both the brokers
            Field ranking = ((SimpleLoadManagerImpl) pulsarServices[i].getLoadManager().get()).getClass()
                    .getDeclaredField("sortedRankings");
            ranking.setAccessible(true);
            AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking = (AtomicReference<Map<Long, Set<ResourceUnit>>>) ranking
                    .get(pulsarServices[i].getLoadManager().get());
            printSortedRanking(sortedRanking);

            // all brokers have same rank to it would be 0 --> set-of-all-the-brokers
            int brokerCount = 0;
            for (Map.Entry<Long, Set<ResourceUnit>> entry : sortedRanking.get().entrySet()) {
                brokerCount += entry.getValue().size();
            }
            assertEquals(brokerCount, BROKER_COUNT);
            TopicName topicName = TopicName.get("persistent://pulsar/use/primary-ns/test-topic");
            ResourceUnit found = pulsarServices[i].getLoadManager().get()
                    .getLeastLoaded(pulsarServices[i].getNamespaceService().getBundle(topicName)).get();
            assertNotNull(found);
        }
    }

    /*
     * tests rankings get updated when we write write the new load reports to the zookeeper on loadbalance root node
     * tests writing pre-configured load report on the zookeeper translates the pre-calculated rankings
     */
    @Test
    public void testUpdateLoadReportAndCheckUpdatedRanking() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            LoadReport lr = new LoadReport();
            lr.setName(lookupAddresses[i]);
            SystemResourceUsage sru = new SystemResourceUsage();
            sru.setBandwidthIn(new ResourceUsage(256, 1024000));
            sru.setBandwidthOut(new ResourceUsage(250, 1024000));
            sru.setMemory(new ResourceUsage(1024, 8192));
            sru.setCpu(new ResourceUsage(5, 400));
            lr.setSystemResourceUsage(sru);

            Whitebox.setInternalState(pulsarServices[0].getLoadManager().get(), "lastLoadReport", lr);
            ResourceLock<LoadReport> lock = Whitebox.getInternalState(pulsarServices[i].getLoadManager().get(),
                    "brokerLock");
            lock.updateValue(lr).join();
        }

        for (int i = 0; i < BROKER_COUNT; i++) {
            Method updateRanking = Whitebox.getMethod(SimpleLoadManagerImpl.class, "updateRanking");
            updateRanking.invoke(pulsarServices[0].getLoadManager().get());
        }

        // do lookup for bunch of bundles
        int totalNamespaces = 200;
        Map<String, Integer> namespaceOwner = new HashMap<>();
        for (int i = 0; i < totalNamespaces; i++) {
            TopicName topicName = TopicName.get("persistent://pulsar/use/primary-ns-" + i + "/test-topic");
            ResourceUnit found = pulsarServices[0].getLoadManager().get()
                    .getLeastLoaded(pulsarServices[0].getNamespaceService().getBundle(topicName)).get();
            if (namespaceOwner.containsKey(found.getResourceId())) {
                namespaceOwner.put(found.getResourceId(), namespaceOwner.get(found.getResourceId()) + 1);
            } else {
                namespaceOwner.put(found.getResourceId(), 1);
            }
        }
        // assert that distribution variation is not more than 20%
        int averageNamespaces = totalNamespaces / BROKER_COUNT;
        int tenPercentOfAverageNamespaces = averageNamespaces / 10;
        int lowerBound = averageNamespaces - tenPercentOfAverageNamespaces;
        int upperBound = averageNamespaces + tenPercentOfAverageNamespaces;

        // assert each broker received ownership of fair amount of namespaces 90%+
        for (Map.Entry<String, Integer> broker : namespaceOwner.entrySet()) {
            log.info("Count of bundles assigned: {}, {} -- lower-bound: {} - upper-bound: {} ",
                    broker.getKey(), broker.getValue(), lowerBound, upperBound);
            assertTrue(broker.getValue() >= lowerBound && broker.getValue() <= upperBound);
        }
    }

    private AtomicReference<Map<Long, Set<ResourceUnit>>> getSortedRanking(PulsarService pulsar)
            throws NoSuchFieldException, IllegalAccessException {
        Field ranking = ((SimpleLoadManagerImpl) pulsar.getLoadManager().get()).getClass()
                .getDeclaredField("sortedRankings");
        ranking.setAccessible(true);
        @SuppressWarnings("unchecked")
        AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking = (AtomicReference<Map<Long, Set<ResourceUnit>>>) ranking
                .get(pulsar.getLoadManager().get());
        return sortedRanking;
    }

    private void printSortedRanking(AtomicReference<Map<Long, Set<ResourceUnit>>> sortedRanking) {
        log.info("Sorted Ranking Result:");
        sortedRanking.get().forEach((score, rus) -> {
            for (ResourceUnit ru : rus) {
                log.info("  - {}, {}", ru.getResourceId(), score);
            }
        });
    }

    /*
     * Pre-publish load report to ZK, each broker has: - Difference memory capacity, for the first 3 brokers memory is
     * bottleneck, for the 4/5th brokers CPU become bottleneck since memory is big enough - non-bundles assigned so all
     * idle resources are available for new bundle Check the broker rankings are the load percentage of each broker.
     */
    @Test(timeOut = 30000)
    public void testBrokerRanking() throws Exception {
        for (int i = 0; i < BROKER_COUNT; i++) {
            LoadReport lr = new LoadReport();
            lr.setName(lookupAddresses[i]);
            SystemResourceUsage sru = new SystemResourceUsage();
            sru.setBandwidthIn(new ResourceUsage(0, 1024000));
            sru.setBandwidthOut(new ResourceUsage(0, 1024000));
            sru.setMemory(new ResourceUsage(1024, 2048 * (i + 1)));
            sru.setCpu(new ResourceUsage(60, 400));
            lr.setSystemResourceUsage(sru);

            ResourceLock<LoadReport> lock = Whitebox.getInternalState(pulsarServices[i].getLoadManager().get(),
                    "brokerLock");
            lock.updateValue(lr).join();
        }

        for (int i = 0; i < BROKER_COUNT; i++) {
            Method updateRanking = Whitebox.getMethod(SimpleLoadManagerImpl.class, "updateRanking");
            updateRanking.invoke(pulsarServices[i].getLoadManager().get());
        }

        // check the ranking result
{'fraction_non_alphanumeric': 0.10540751129811439, 'fraction_numerical': 0.008508648901355774, 'mean_word_length': 3.105694177863084, 'pattern_counts': {'":': 8, '<': 77, '<?xml version=': 0, '>': 71, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 4, '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': '7063625', 'n_tokens_mistral': 9230, 'n_tokens_neox': 8884, 'n_words': 1980}