From 4abffa9f2461918a5858650115420b21591ae03e Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Sat, 8 Feb 2014 09:33:05 +1100
Subject: [PATCH] Optimize / clean up server disconnect method

---
 .../net/md_5/bungee/ServerConnection.java     | 19 +++----------------
 .../java/net/md_5/bungee/ServerConnector.java |  2 --
 .../net/md_5/bungee/netty/ChannelWrapper.java | 18 +++++++++++++++---
 3 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnection.java b/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
index 3fd2e982..61a8b841 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnection.java
@@ -1,7 +1,7 @@
 package net.md_5.bungee;
 
+import io.netty.channel.ChannelFutureListener;
 import java.net.InetSocketAddress;
-import java.util.concurrent.TimeUnit;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.Setter;
@@ -41,7 +41,7 @@ public class ServerConnection implements Server
     }
 
     @Override
-    public synchronized void disconnect(String reason)
+    public void disconnect(String reason)
     {
         disconnect( TextComponent.fromLegacyText( reason ) );
     }
@@ -49,20 +49,7 @@ public class ServerConnection implements Server
     @Override
     public void disconnect(BaseComponent... reason)
     {
-        if ( !ch.isClosed() )
-        {
-            // TODO: Can we just use a future here?
-            unsafe().sendPacket( new Kick( ComponentSerializer.toString( reason ) ) );
-            ch.getHandle().eventLoop().schedule( new Runnable()
-            {
-                @Override
-                public void run()
-                {
-                    ch.getHandle().close();
-                }
-            }, 100, TimeUnit.MILLISECONDS );
-        }
-
+        ch.write( new Kick( ComponentSerializer.toString( reason ) ), ChannelFutureListener.CLOSE );
     }
 
     @Override
diff --git a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
index 4d5f9c84..d8981ad2 100644
--- a/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
+++ b/proxy/src/main/java/net/md_5/bungee/ServerConnector.java
@@ -1,8 +1,6 @@
 package net.md_5.bungee;
 
 import com.google.common.base.Preconditions;
-import com.google.common.io.ByteArrayDataOutput;
-import com.google.common.io.ByteStreams;
 import java.util.Objects;
 import java.util.Queue;
 import lombok.RequiredArgsConstructor;
diff --git a/proxy/src/main/java/net/md_5/bungee/netty/ChannelWrapper.java b/proxy/src/main/java/net/md_5/bungee/netty/ChannelWrapper.java
index 93962e27..d994ff63 100644
--- a/proxy/src/main/java/net/md_5/bungee/netty/ChannelWrapper.java
+++ b/proxy/src/main/java/net/md_5/bungee/netty/ChannelWrapper.java
@@ -3,6 +3,8 @@ package net.md_5.bungee.netty;
 import net.md_5.bungee.protocol.PacketWrapper;
 import com.google.common.base.Preconditions;
 import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import lombok.Getter;
@@ -34,22 +36,32 @@ public class ChannelWrapper
         ch.pipeline().get( MinecraftEncoder.class ).setProtocolVersion( protocol );
     }
 
-    public synchronized void write(Object packet)
+    public synchronized void write(Object packet, ChannelFutureListener future)
     {
         if ( !closed )
         {
+            ChannelFuture listener;
             if ( packet instanceof PacketWrapper )
             {
                 ( (PacketWrapper) packet ).setReleased( true );
-                ch.write( ( (PacketWrapper) packet ).buf, ch.voidPromise() );
+                listener = ch.write( ( (PacketWrapper) packet ).buf, ch.voidPromise() );
             } else
             {
-                ch.write( packet, ch.voidPromise() );
+                listener = ch.write( packet, ch.voidPromise() );
+            }
+            if ( future != null )
+            {
+                listener.addListener( future );
             }
             ch.flush();
         }
     }
 
+    public synchronized void write(Object packet)
+    {
+        write( packet, null );
+    }
+
     public synchronized void close()
     {
         if ( !closed )