Add kuvel loadbalancer register/unregister subcommands

Agent-Logs-Url: https://github.com/AzisabaNetwork/Kuvel/sessions/f0c88feb-5760-41f0-80c1-b94209a0e7b2

Co-authored-by: matsu1213 <85377530+matsu1213@users.noreply.github.com>

Author: Copilot

src/main/java/net/azisaba/kuvel/KuvelServiceHandler.java

@@ -5,6 +5,8 @@
 import io.fabric8.kubernetes.api.model.ContainerPort;
 import io.fabric8.kubernetes.api.model.Pod;
 import io.fabric8.kubernetes.api.model.PodList;
+import io.fabric8.kubernetes.api.model.apps.ReplicaSet;
+import io.fabric8.kubernetes.api.model.apps.ReplicaSetList;
 import io.fabric8.kubernetes.client.KubernetesClient;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
@@ -17,11 +19,13 @@
 
 import io.fabric8.kubernetes.client.dsl.FilterWatchListDeletable;
 import io.fabric8.kubernetes.client.dsl.PodResource;
+import io.fabric8.kubernetes.client.dsl.RollableScalableResource;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import net.azisaba.kuvel.discovery.LoadBalancerDiscovery;
 import net.azisaba.kuvel.discovery.ServerDiscovery;
 import net.azisaba.kuvel.loadbalancer.LoadBalancer;
+import net.azisaba.kuvel.loadbalancer.strategy.impl.RoundRobinLoadBalancingStrategy;
 import net.azisaba.kuvel.util.LabelKeys;
 import net.azisaba.kuvel.util.UidAndServerNameMap;
 
@@ -337,4 +341,63 @@ public void unregisterPod(Pod pod) {
   public boolean isPodRegistered(String podId) {
     return podUidAndServerNameMap.getServerNameFromUid(podId) != null;
   }
+
+  /**
+   * Register a load balancer with ReplicaSet uid and server name.
+   *
+   * @param replicaSetUid The ReplicaSet uid to register.
+   * @param serverName The server name to register.
+   * @return true if the load balancer is registered successfully.
+   */
+  public boolean registerLoadBalancer(String replicaSetUid, String serverName) {
+    FilterWatchListDeletable<ReplicaSet, ReplicaSetList, RollableScalableResource<ReplicaSet>> request =
+        client.apps().replicaSets().inNamespace(namespace);
+
+    for (Entry<String, String> e : plugin.getKuvelConfig().getLabelSelectors().entrySet()) {
+      request = request.withLabel(e.getKey(), e.getValue());
+    }
+
+    Optional<ReplicaSet> replicaSet =
+        request
+            .list()
+            .getItems()
+            .stream()
+            .filter(r -> r.getMetadata().getUid().equals(replicaSetUid))
+            .findFirst();
+    if (replicaSet.isEmpty()) {
+      return false;
+    }
+
+    boolean initialServer =
+        replicaSet
+            .get()
+            .getMetadata()
+            .getLabels()
+            .getOrDefault(
+                LabelKeys.INITIAL_SERVER.getKey(plugin.getKuvelConfig().getLabelKeyPrefix()),
+                "false")
+            .equalsIgnoreCase("true");
+
+    plugin
+        .getProxy()
+        .registerServer(new ServerInfo(serverName, new InetSocketAddress("0.0.0.0", 0)));
+    registerLoadBalancer(
+        new LoadBalancer(
+            plugin.getProxy(),
+            plugin.getProxy().getServer(serverName).orElseThrow(),
+            new RoundRobinLoadBalancingStrategy(),
+            replicaSetUid,
+            initialServer));
+    return true;
+  }
+
+  /**
+   * Gets whether the specified load balancer is registered.
+   *
+   * @param replicaSetUid The ReplicaSet uid to check.
+   * @return true if the specified load balancer uid is registered.
+   */
+  public boolean isLoadBalancerRegistered(String replicaSetUid) {
+    return replicaSetUidAndServerNameMap.getServerNameFromUid(replicaSetUid) != null;
+  }
 }

src/main/java/net/azisaba/kuvel/command/KuvelCommand.java

@@ -93,6 +93,32 @@ public void execute(Invocation invocation) {
     }
 
     if (args[0].equalsIgnoreCase("register")) {
+      if (args.length >= 2 && args[1].equalsIgnoreCase("loadbalancer")) {
+        if (args.length < 4) {
+          source.sendMessage(
+              Component.text("Usage: /kuvel register loadbalancer <replicaSetUid> <serverName>"));
+          return;
+        }
+        String mappedUid = handler.getReplicaSetUidAndServerNameMap().getUidFromServerName(args[3]);
+        if (mappedUid != null && !mappedUid.equals(args[2])) {
+          source.sendMessage(
+              Component.text("Server name is already mapped to another load balancer uid: " + mappedUid));
+          return;
+        }
+        if (plugin.getProxy().getServer(args[3]).isPresent() && mappedUid == null) {
+          source.sendMessage(Component.text("Server name is already registered: " + args[3]));
+          return;
+        }
+
+        if (handler.registerLoadBalancer(args[2], args[3])) {
+          source.sendMessage(
+              Component.text("Registered load balancer mapping: " + args[2] + " -> " + args[3]));
+        } else {
+          source.sendMessage(Component.text("Failed to register load balancer mapping."));
+        }
+        return;
+      }
+
       if (args.length < 3) {
         source.sendMessage(Component.text("Usage: /kuvel register <podUid> <serverName>"));
         return;
@@ -117,6 +143,23 @@ public void execute(Invocation invocation) {
     }
 
     if (args[0].equalsIgnoreCase("unregister")) {
+      if (args.length >= 2 && args[1].equalsIgnoreCase("loadbalancer")) {
+        if (args.length < 3) {
+          source.sendMessage(Component.text("Usage: /kuvel unregister loadbalancer <replicaSetUid>"));
+          return;
+        }
+
+        if (!handler.isLoadBalancerRegistered(args[2])) {
+          source.sendMessage(
+              Component.text("Load balancer uid is not currently registered: " + args[2]));
+          return;
+        }
+
+        handler.unregisterLoadBalancer(args[2]);
+        source.sendMessage(Component.text("Unregistered load balancer mapping: " + args[2]));
+        return;
+      }
+
       if (args.length < 2) {
         source.sendMessage(Component.text("Usage: /kuvel unregister <podUid>"));
         return;
@@ -187,13 +230,19 @@ public void execute(Invocation invocation) {
       }
 
       int cleanedLoadBalancerCount = 0;
+      int repairedLoadBalancerCount = 0;
       for (Map.Entry<String, String> entry :
           handler.getReplicaSetUidAndServerNameMap().getAllMap().entrySet()) {
         if (plugin.getProxy().getServer(entry.getValue()).isPresent()) {
           continue;
         }
-        handler.unregisterLoadBalancer(entry.getKey());
-        cleanedLoadBalancerCount++;
+        boolean success = handler.registerLoadBalancer(entry.getKey(), entry.getValue());
+        if (success && plugin.getProxy().getServer(entry.getValue()).isPresent()) {
+          repairedLoadBalancerCount++;
+        } else {
+          handler.unregisterLoadBalancer(entry.getKey());
+          cleanedLoadBalancerCount++;
+        }
       }
 
       source.sendMessage(
@@ -202,6 +251,8 @@ public void execute(Invocation invocation) {
                   + repairedPodCount
                   + ", cleanedPod="
                   + cleanedPodCount
+                  + ", repairedLoadBalancer="
+                  + repairedLoadBalancerCount
                   + ", cleanedLoadBalancer="
                   + cleanedLoadBalancerCount));
       return;
@@ -214,7 +265,9 @@ private void sendHelp(CommandSource source) {
     source.sendMessage(Component.text("/kuvel status"));
     source.sendMessage(Component.text("/kuvel list <pods|loadbalancers>"));
     source.sendMessage(Component.text("/kuvel register <podUid> <serverName>"));
+    source.sendMessage(Component.text("/kuvel register loadbalancer <replicaSetUid> <serverName>"));
     source.sendMessage(Component.text("/kuvel unregister <podUid>"));
+    source.sendMessage(Component.text("/kuvel unregister loadbalancer <replicaSetUid>"));
     source.sendMessage(Component.text("/kuvel setname <podUid> <serverName>"));
     source.sendMessage(Component.text("/kuvel repair"));
   }