Multi-raft-group setup rebalance without PD

jraft-core/src/main/java/com/alipay/sofa/jraft/CliService.java

@@ -17,6 +17,8 @@
 package com.alipay.sofa.jraft;
 
 import java.util.List;
+import java.util.Map;
+import java.util.Queue;
 
 import com.alipay.sofa.jraft.conf.Configuration;
 import com.alipay.sofa.jraft.entity.PeerId;
@@ -118,4 +120,14 @@ public interface CliService extends Lifecycle<CliOptions> {
      * @return all alive peers of the replication group
      */
     List<PeerId> getAlivePeers(final String groupId, final Configuration conf);
+
+    /**
+     * Balance the number of leaders.
+     *
+     * @param groupIds  the raft group id queue
+     * @param conf      current configuration
+     * @param leaderIds the ids of raft group leader
+     * @return operation status
+     */
+    Status rebalance(final List<String> groupIds, final Configuration conf, final Map<String, PeerId> leaderIds);
 }

jraft-core/src/main/java/com/alipay/sofa/jraft/core/CliServiceImpl.java

@@ -16,8 +16,7 @@
  */
 package com.alipay.sofa.jraft.core;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
@@ -330,7 +329,7 @@ public Status getLeader(final String groupId, final Configuration conf, final Pe
         final Status st = new Status(-1, "Fail to get leader of group %s", groupId);
         for (final PeerId peer : conf) {
             if (!this.cliClientService.connect(peer.getEndpoint())) {
-                LOG.error("Fail to connect peer {} to get leader for group {}.", groupId);
+                LOG.error("Fail to connect peer {} to get leader for group {}.", peer, groupId);
                 continue;
             }
 
@@ -383,6 +382,66 @@ public List<PeerId> getAlivePeers(final String groupId, final Configuration conf
         return getPeers(groupId, conf, true);
     }
 
+    @Override
+    public Status rebalance(final List<String> groupIds, final Configuration conf, final Map<String, PeerId> leaderIds) {
+        Requires.requireTrue(!groupIds.isEmpty(), "Empty group id list");
+        Requires.requireNonNull(conf, "Null configuration");
+        Requires.requireTrue(!conf.isEmpty(), "No peers of configuration");
+
+        final int groupSizePerPeer = groupIds.size() / conf.size();
+        final Queue<String> groupDeque = new ArrayDeque<>(groupIds);
+        final Map<PeerId, Integer> peerMap = new HashMap<>();
+        for (;;) {
+            final String groupId = groupDeque.poll();
+            if (StringUtils.isEmpty(groupId)) {
+                break;
+            }
+            final PeerId leaderId = new PeerId();
+            try {
+                final Status status = getLeader(groupId, conf, leaderId);
+                if (!status.isOk()) {
+                    throw new JRaftException("No leader in group: " + groupId);
+                }
+                if (leaderId.getEndpoint() == null) {
+                    continue;
+                }
+                LOG.info("Group {} leader is {}.", groupId, leaderId);
+            } catch (final Exception e) {
+                groupDeque.add(groupId);
+                continue;
+            }
+            final Integer size = peerMap.get(leaderId);
+            if (size == null) {
+                peerMap.put(leaderId, 1);
+                continue;
+            }
+            if (size <= groupSizePerPeer) {
+                peerMap.put(leaderId, size + 1);
+                continue;
+            }
+            for (final PeerId peerId : getAlivePeers(groupId, conf)) {
+                final Integer pSize = peerMap.get(peerId);
+                if (pSize != null && pSize >= groupSizePerPeer) {
+                    continue;
+                }
+                try {
+                    final Status status = transferLeader(groupId, conf, peerId);
+                    if (status.isOk()) {
+                        LOG.info("Group {} transfer leader to {}.", groupId, peerId);
+                        groupDeque.add(groupId);
+                        break;
+                    } else {
+                        LOG.error("Fail to transfer leader to {}.", peerId);
+                    }
+                } catch (final Exception e) {
+                    LOG.error("Fail to transfer leader to {}.", peerId);
+                }
+            }
+            leaderIds.put(groupId, leaderId);
+        }
+        return Status.OK();
+    }
+
     private List<PeerId> getPeers(final String groupId, final Configuration conf, final boolean onlyGetAlive) {
         Requires.requireTrue(!StringUtils.isBlank(groupId), "Blank group id");
         Requires.requireNonNull(conf, "Null conf");

jraft-core/src/test/java/com/alipay/sofa/jraft/core/CliServiceTest.java

@@ -18,26 +18,24 @@
 
 import java.io.File;
 import java.nio.ByteBuffer;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import com.alipay.sofa.jraft.*;
 import org.apache.commons.io.FileUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
-import com.alipay.sofa.jraft.CliService;
-import com.alipay.sofa.jraft.Node;
-import com.alipay.sofa.jraft.NodeManager;
-import com.alipay.sofa.jraft.RouteTable;
 import com.alipay.sofa.jraft.conf.Configuration;
 import com.alipay.sofa.jraft.entity.PeerId;
 import com.alipay.sofa.jraft.entity.Task;
 import com.alipay.sofa.jraft.option.CliOptions;
 import com.alipay.sofa.jraft.test.TestUtils;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.mockito.Spy;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -51,14 +49,17 @@ public class CliServiceTest {
     private String        dataPath;
 
     private TestCluster   cluster;
-    private final String  groupId = "CliServiceTest";
 
-    private CliService    cliService;
+    private final String  groupId    = "CliServiceTest";
+
+    @Spy
+    private CliService    cliService = new CliServiceImpl();
 
     private Configuration conf;
 
     @Before
     public void setup() throws Exception {
+        MockitoAnnotations.initMocks(this);
 
         this.dataPath = TestUtils.mkTempDir();
         FileUtils.forceMkdir(new File(this.dataPath));
@@ -71,7 +72,6 @@ public void setup() throws Exception {
         }
         cluster.waitLeader();
 
-        cliService = new CliServiceImpl();
         this.conf = new Configuration(peers);
         assertTrue(cliService.init(new CliOptions()));
     }
@@ -242,4 +242,24 @@ public void testGetAlivePeers() throws Exception {
             assertEquals("Fail to get leader of group " + this.groupId, e.getMessage());
         }
     }
+
+    @Test
+    public void testRebalance() {
+        final PeerId leader = cluster.getLeader().getNodeId().getPeerId().copy();
+        assertNotNull(leader);
+
+        Mockito.doReturn(this.conf.getPeers()).when(this.cliService).getAlivePeers(groupId, this.conf);
+        for (final PeerId peer : this.conf.getPeerSet()) {
+            Mockito.doReturn(Status.OK()).when(this.cliService).transferLeader(groupId, this.conf, peer);
+        }
+
+        final List<String> groupIds = new ArrayList<>();
+        groupIds.add(groupId);
+        groupIds.add(groupId);
+        groupIds.add(groupId);
+
+        final Map<String, PeerId> leaderIds = new HashMap<>();
+        assertTrue(this.cliService.rebalance(groupIds, this.conf, leaderIds).isOk());
+        assertEquals(1, leaderIds.size());
+    }
 }