Improve memory management

Rewrite timer tasks as non-inner class. It works around problem with
garbage collecting BGPFsm.

3 files changed, 80 insertions, 32 deletions

diff --git a/src/main/java/com/lumaserv/bgp/BGPFsm.java b/src/main/java/com/lumaserv/bgp/BGPFsm.java
index 533cb96..7c5b6aa 100644
--- a/src/main/java/com/lumaserv/bgp/BGPFsm.java
+++ b/src/main/java/com/lumaserv/bgp/BGPFsm.java
@@ -12,7 +12,7 @@ import java.util.Optional;
 import java.util.Timer;
 import java.util.TimerTask;
 
-public class BGPFsm {
+public class BGPFsm implements Callback {
     public static final int CONFIG_HOLD_TIME = 80;  // FIXME use from configuration
     public static final byte CONFIG_VERSION = 4;
 
@@ -502,7 +502,8 @@ public class BGPFsm {
         public String toString() { return "Established"; }
     }
 
-    BGPFsm() {
+    BGPFsm(BGPSession session) {
+        this.session = session;
         IDLE = new Idle();
         CONNECT = new Connect();
         ACTIVE = new Active();
@@ -515,6 +516,27 @@ public class BGPFsm {
         connectRetryTimer = new Timer();
     }
 
+    public void cleanup() {
+        System.out.println("BGPFsm cleanup");
+        if (keepAliveTimerTask != null) {
+            keepAliveTimerTask.cancel();
+            keepAliveTimerTask = null;
+        }
+        keepAliveTimer = null;
+
+        if (holdTimerTask != null) {
+            holdTimerTask.cancel();
+            holdTimerTask = null;
+        }
+        holdTimer = null;
+
+        if (connectRetryTimerTask != null) {
+            connectRetryTimerTask.cancel();
+            connectRetryTimerTask = null;
+        }
+        connectRetryTimer = null;
+    }
+
     void setKeepAliveTimer() {
         //        System.out.println("setKeepAliveTimer: " + holdTime / 3);
 
@@ -526,11 +548,7 @@ public class BGPFsm {
         if (holdTime <= 0)
             return;
 
-        keepAliveTimerTask = new TimerTask() {
-                public void run() {
-                    currentState.keepAliveTimerExpires();
-                }
-            };
+        keepAliveTimerTask = new FsmTimerTask(this);
         keepAliveTimer.schedule(keepAliveTimerTask, 1000 * holdTime / 3);
     }
 
@@ -544,11 +562,7 @@ public class BGPFsm {
         if (delay <= 0)
             return;
 
-        holdTimerTask = new TimerTask() {
-                public void run() {
-                    currentState.holdTimerExpires();
-                }
-            };
+        holdTimerTask = new FsmTimerTask(this);
         holdTimer.schedule(holdTimerTask, 1000 * delay);
     }
 
@@ -562,11 +576,44 @@ public class BGPFsm {
         if (delay <= 0)
             return;
 
-        connectRetryTimerTask = new TimerTask() {
-                public void run() {
-                    currentState.connectRetryTimerExpires();
-                }
-            };
+        connectRetryTimerTask = new FsmTimerTask(this);
         connectRetryTimer.schedule(connectRetryTimerTask, 1000 * delay);
     }
+
+    public void callback(Object task) {
+        if (task == connectRetryTimerTask) {
+            currentState.connectRetryTimerExpires();
+        } else if (task == holdTimerTask) {
+            currentState.holdTimerExpires();
+        } else if (task == keepAliveTimerTask) {
+            currentState.keepAliveTimerExpires();
+        } else {
+            throw new RuntimeException("Unknown timer task: " + task);
+        }
+    }
+
+    static class FsmTimerTask extends TimerTask {
+        Callback owner;
+
+        FsmTimerTask(Callback owner) {
+            this.owner = owner;
+        }
+
+        public void run() {
+            if (owner != null) {
+                owner.callback(this);
+            }
+        }
+
+        public boolean cancel() {
+            boolean res = super.cancel();
+            // Break cyclic references to allow GC
+            this.owner = null;
+            return res;
+        }
+    }
+}
+
+interface Callback {
+    void callback(Object cb);
 }
diff --git a/src/main/java/com/lumaserv/bgp/BGPServer.java b/src/main/java/com/lumaserv/bgp/BGPServer.java
index e95caea..78ceee8 100644
--- a/src/main/java/com/lumaserv/bgp/BGPServer.java
+++ b/src/main/java/com/lumaserv/bgp/BGPServer.java
@@ -99,6 +99,7 @@ public class BGPServer implements Runnable {
                 }
                 sessions.remove(otherSession);
                 otherSession.dropConnection();
+                otherSession.cleanup();
                 // Accept new connection
                 return false;
             } else {
@@ -128,12 +129,9 @@ public class BGPServer implements Runnable {
                     continue;
                 }
 
-                BGPFsm fsm = new BGPFsm();
-                BGPSession session = new BGPSession(this, socket, config, fsm);
-                fsm.setSession(session);
+                BGPSession session = new BGPSession(this, socket, config);
                 sessions.add(session);
-                System.out.println("automaticStartPassive");
-                fsm.getCurrentState().automaticStartPassive();
+                session.getFsm().getCurrentState().automaticStartPassive();
 
                 // System.out.println("tcpConnectionConfirmed in " + fsm.getCurrentState());
                 // fsm.getCurrentState().tcpConnectionConfirmed();
@@ -158,11 +156,9 @@ public class BGPServer implements Runnable {
         }
 
         try {
-            BGPFsm fsm = new BGPFsm();
-            BGPSession session = new BGPSession(this, config, fsm, host, port);
-            fsm.setSession(session);
+            BGPSession session = new BGPSession(this, config, host, port);
             sessions.add(session);
-            fsm.getCurrentState().automaticStart();
+            session.getFsm().getCurrentState().automaticStart();
             return true;
         } catch (IOException ex) {
             ex.printStackTrace();
@@ -173,6 +169,7 @@ public class BGPServer implements Runnable {
     public void shutdown() {
         for(BGPSession session : sessions) {
             session.automaticStop();
+            session.cleanup();
         }
         sessions.clear();
         if (serverSocket != null) {
diff --git a/src/main/java/com/lumaserv/bgp/BGPSession.java b/src/main/java/com/lumaserv/bgp/BGPSession.java
index 9de32ab..d51a83e 100644
--- a/src/main/java/com/lumaserv/bgp/BGPSession.java
+++ b/src/main/java/com/lumaserv/bgp/BGPSession.java
@@ -39,21 +39,21 @@ public class BGPSession implements Runnable {
     @Getter
     byte[] remoteIdentifier;
 
-    public BGPSession(BGPServer server, Socket socket, BGPSessionConfiguration configuration, BGPFsm fsm) throws IOException {
+    public BGPSession(BGPServer server, Socket socket, BGPSessionConfiguration configuration) throws IOException {
         this.server = server;
         this.configuration = configuration;
         this.socket = socket;
         this.inputStream = socket.getInputStream();
         this.outputStream = socket.getOutputStream();
-        this.fsm = fsm;
+        this.fsm = new BGPFsm(this);
         this.host = null;
         this.port = -1;
     }
 
-    public BGPSession(BGPServer server, BGPSessionConfiguration configuration, BGPFsm fsm, String host, int port) throws IOException {
+    public BGPSession(BGPServer server, BGPSessionConfiguration configuration, String host, int port) throws IOException {
         this.server = server;
         this.configuration = configuration;
-        this.fsm = fsm;
+        this.fsm = new BGPFsm(this);
         this.host = host;
         this.port = port;
     }
@@ -68,6 +68,11 @@ public class BGPSession implements Runnable {
 
     protected void finalize() throws Throwable {
         System.out.println("Finalize BGPSession");
+        super.finalize();
+    }
+
+    public void cleanup() {
+        fsm.cleanup();
     }
 
     public void automaticStop() {
@@ -106,7 +111,7 @@ public class BGPSession implements Runnable {
                     } catch(IOException ex) {
                     }
                     dropConnection();
-                    // TODO Break links between fsm and session!
+                    cleanup();
                 } else {
                     isAS4Capability = isAdvAS4Capability && open.getAS4Capability().isPresent();
                     remoteIdentifier = open.getIdentifier();
@@ -224,7 +229,6 @@ System.out.println("Sent open");
         } catch (IOException ex) {
             if (!closed) {
                 fsm.getCurrentState().tcpConnectionFails();
-                ex.printStackTrace();
             }
             closed = true;
             configuration.getListener().onClose(this);