From 12ad0c792be29020cedd3af29216d2250e037d25 Mon Sep 17 00:00:00 2001
From: Tux <write@imaginarycode.com>
Date: Sun, 6 Dec 2015 16:31:21 -0500
Subject: [PATCH] One more attempt to fix player count. Now it's calculated
 with Lua.

---
 .../minecraft/redisbungee/RedisBungee.java    | 54 ++++++-------------
 src/main/resources/lua/get_player_count.lua   | 24 +++++++++
 2 files changed, 39 insertions(+), 39 deletions(-)
 create mode 100644 src/main/resources/lua/get_player_count.lua

diff --git a/src/main/java/com/imaginarycode/minecraft/redisbungee/RedisBungee.java b/src/main/java/com/imaginarycode/minecraft/redisbungee/RedisBungee.java
index 5efd72c..ac2b4e7 100644
--- a/src/main/java/com/imaginarycode/minecraft/redisbungee/RedisBungee.java
+++ b/src/main/java/com/imaginarycode/minecraft/redisbungee/RedisBungee.java
@@ -20,7 +20,6 @@ import net.md_5.bungee.config.ConfigurationProvider;
 import net.md_5.bungee.config.YamlConfiguration;
 import redis.clients.jedis.*;
 import redis.clients.jedis.exceptions.JedisConnectionException;
-import redis.clients.jedis.exceptions.JedisException;
 
 import java.io.*;
 import java.util.*;
@@ -60,6 +59,7 @@ public final class RedisBungee extends Plugin {
     private ScheduledTask heartbeatTask;
     private boolean usingLua;
     private LuaManager.Script serverToPlayersScript;
+    private LuaManager.Script getPlayerCountScript;
 
     /**
      * Fetch the {@link RedisBungeeAPI} object created on plugin start.
@@ -121,51 +121,26 @@ public final class RedisBungee extends Plugin {
     }
 
     final Multimap<String, UUID> serversToPlayers() {
-        if (usingLua) {
-            Collection<String> data = (Collection<String>) serverToPlayersScript.eval(ImmutableList.<String>of(), getServerIds());
+        Collection<String> data = (Collection<String>) serverToPlayersScript.eval(ImmutableList.<String>of(), getServerIds());
 
-            ImmutableMultimap.Builder<String, UUID> builder = ImmutableMultimap.builder();
-
-            // TODO: This seems pretty slow, but execution times over the long term seem to stay below that of the
-            // Java implementation, at least. If you have a better idea, I want to see it!
-            String key = null;
-
-            for (String s : data) {
-                if (key == null) {
-                    key = s;
-                    continue;
-                }
-
-                builder.put(key, UUID.fromString(s));
-                key = null;
+        ImmutableMultimap.Builder<String, UUID> builder = ImmutableMultimap.builder();
+        String key = null;
+        for (String s : data) {
+            if (key == null) {
+                key = s;
+                continue;
             }
 
-            return builder.build();
-        } else {
-            ImmutableMultimap.Builder<String, UUID> multimapBuilder = ImmutableMultimap.builder();
-            for (UUID p : getPlayers()) {
-                String name = dataManager.getServer(p);
-                if (name != null)
-                    multimapBuilder.put(name, p);
-            }
-            return multimapBuilder.build();
+            builder.put(key, UUID.fromString(s));
+            key = null;
         }
+
+        return builder.build();
     }
 
     final int getCount() {
-        int c = 0;
-        if (pool != null) {
-            try (Jedis rsc = pool.getResource()) {
-                for (String i : getServerIds()) {
-                    c += rsc.scard("proxy:" + i + ":usersOnline");
-                }
-            } catch (JedisConnectionException e) {
-                // Redis server has disappeared!
-                getLogger().log(Level.SEVERE, "Unable to get connection from pool - did your Redis server go away?", e);
-                throw new RuntimeException("Unable to get total player count", e);
-            }
-        }
-        return c;
+        Long count = (Long) getPlayerCountScript.eval(ImmutableList.<String>of(), ImmutableList.<String>of());
+        return count.intValue();
     }
 
     private Set<String> getLocalPlayersAsUuidStrings() {
@@ -250,6 +225,7 @@ public final class RedisBungee extends Plugin {
                         } else {
                             LuaManager manager = new LuaManager(this);
                             serverToPlayersScript = manager.createScript(IOUtil.readInputStreamAsString(getResourceAsStream("lua/server_to_players.lua")));
+                            getPlayerCountScript = manager.createScript(IOUtil.readInputStreamAsString(getResourceAsStream("lua/get_player_count.lua")));
                         }
                         break;
                     }
diff --git a/src/main/resources/lua/get_player_count.lua b/src/main/resources/lua/get_player_count.lua
new file mode 100644
index 0000000..d337320
--- /dev/null
+++ b/src/main/resources/lua/get_player_count.lua
@@ -0,0 +1,24 @@
+local c = redis.call
+
+local curTime = c("TIME")
+local time = tonumber(curTime[1])
+
+local heartbeats = c("HGETALL", "heartbeats")
+local total = 0
+local key
+
+for _, v in pairs(heartbeats) do
+    if not key then
+        key = v
+    else
+        local n = tonumber(v)
+        if n then
+            if n + 30 >= time then
+                total = total + c("SCARD", "proxy:" .. key .. ":usersOnline")
+            end
+        end
+        key = nil
+    end
+end
+
+return total