From a778825152973f15a12abece8434e41b62381a18 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 5 Aug 2018 17:10:18 -0400
Subject: [PATCH] Add server registration API.

---
 .../api/proxy/ProxyServer.java                | 27 ++++++++++++
 .../velocitypowered/proxy/VelocityServer.java | 21 ++++++++++
 .../velocitypowered/proxy/util/ServerMap.java | 42 +++++++++++++++----
 3 files changed, 83 insertions(+), 7 deletions(-)

diff --git a/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java b/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
index 7ff2d8ea9..7e23a0a07 100644
--- a/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
+++ b/api/src/main/java/com/velocitypowered/api/proxy/ProxyServer.java
@@ -1,5 +1,7 @@
 package com.velocitypowered.api.proxy;
 
+import com.velocitypowered.api.server.ServerInfo;
+
 import javax.annotation.Nonnull;
 import java.util.Collection;
 import java.util.Optional;
@@ -29,4 +31,29 @@ public interface ProxyServer {
      * @return the players online on this proxy
      */
     Collection<Player> getAllPlayers();
+
+    /**
+     * Retrieves a registered {@link ServerInfo} instance by its name.
+     * @param name the name of the server
+     * @return the server
+     */
+    Optional<ServerInfo> getServerInfo(@Nonnull String name);
+
+    /**
+     * Retrieves all {@link ServerInfo}s registered with this proxy.
+     * @return the servers registered with this proxy
+     */
+    Collection<ServerInfo> getAllServers();
+
+    /**
+     * Registers a server with this proxy. A server with this name should not already exist.
+     * @param server the server to register
+     */
+    void registerServer(@Nonnull ServerInfo server);
+
+    /**
+     * Unregisters this server from the proxy.
+     * @param server the server to unregister
+     */
+    void unregisterServer(@Nonnull ServerInfo server);
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
index a18cd2db8..71659208b 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/VelocityServer.java
@@ -150,4 +150,25 @@ public class VelocityServer implements ProxyServer {
     public Collection<Player> getAllPlayers() {
         return ImmutableList.copyOf(connectionsByUuid.values());
     }
+
+    @Override
+    public Optional<ServerInfo> getServerInfo(@Nonnull String name) {
+        Preconditions.checkNotNull(name, "name");
+        return servers.getServer(name);
+    }
+
+    @Override
+    public Collection<ServerInfo> getAllServers() {
+        return servers.getAllServers();
+    }
+
+    @Override
+    public void registerServer(@Nonnull ServerInfo server) {
+        servers.register(server);
+    }
+
+    @Override
+    public void unregisterServer(@Nonnull ServerInfo server) {
+        servers.unregister(server);
+    }
 }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/util/ServerMap.java b/proxy/src/main/java/com/velocitypowered/proxy/util/ServerMap.java
index eb0948c26..c31b77712 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/util/ServerMap.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/util/ServerMap.java
@@ -8,21 +8,49 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class ServerMap {
     private final Map<String, ServerInfo> servers = new HashMap<>();
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
 
-    public Optional<ServerInfo> getServer(String name) {
-        Preconditions.checkNotNull(name, "name");
-        return Optional.ofNullable(servers.get(name.toLowerCase()));
+    public Optional<ServerInfo> getServer(String server) {
+        Preconditions.checkNotNull(server, "server");
+        lock.readLock().lock();
+        try {
+            return Optional.ofNullable(servers.get(server.toLowerCase()));
+        } finally {
+            lock.readLock().unlock();
+        }
     }
 
     public Collection<ServerInfo> getAllServers() {
-        return ImmutableList.copyOf(servers.values());
+        lock.readLock().lock();
+        try {
+            return ImmutableList.copyOf(servers.values());
+        } finally {
+            lock.readLock().unlock();
+        }
     }
 
-    public void register(ServerInfo info) {
-        Preconditions.checkNotNull(info, "info");
-        servers.put(info.getName(), info);
+    public void register(ServerInfo server) {
+        Preconditions.checkNotNull(server, "server");
+        lock.writeLock().lock();
+        try {
+            Preconditions.checkArgument(servers.putIfAbsent(server.getName(), server) == null, "Server with name %s already registered", server.getName());
+        } finally {
+            lock.writeLock().unlock();
+        }
+    }
+
+    public void unregister(ServerInfo server) {
+        Preconditions.checkNotNull(server, "server");
+        lock.writeLock().lock();
+        try {
+            Preconditions.checkArgument(servers.remove(server.getName(), server), "Server with this name is not registered!");
+        } finally {
+            lock.writeLock().unlock();
+        }
     }
 }