From b1464eba5605768fd88747749e4882ca05088a28 Mon Sep 17 00:00:00 2001
From: noppoman <yuki@miketokyo.com>
Date: Sun, 26 Jun 2016 02:07:45 +0900
Subject: [PATCH] add socketsConnected and closeClientsConnected() to manage
 sockets

---
 Sources/HTTPStream.swift |  7 +++++++
 Sources/Skelton.swift    | 36 ++++++++++++++++++++++++------------
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/Sources/HTTPStream.swift b/Sources/HTTPStream.swift
index 300211b..f52f011 100644
--- a/Sources/HTTPStream.swift
+++ b/Sources/HTTPStream.swift
@@ -25,3 +25,10 @@ extension HTTPStream {
         self.send(chunk, timingOut: deadline, completion: completion)
     }
 }
+
+extension HTTPStream: Equatable {}
+
+public func ==(lhs: HTTPStream, rhs: HTTPStream) -> Bool {
+    return ObjectIdentifier(lhs) == ObjectIdentifier(rhs)
+}
+
diff --git a/Sources/Skelton.swift b/Sources/Skelton.swift
index cec7186..ef37e83 100644
--- a/Sources/Skelton.swift
+++ b/Sources/Skelton.swift
@@ -64,7 +64,11 @@ public final class Skelton {
     private let server: TCPServer
     
     // Current connected clients count.
-    public var clientsConnected: Int = 0
+    public var clientsConnected: Int {
+        return socketsConnected.count
+    }
+    
+    private var socketsConnected = [HTTPStream]()
     
     /**
      - parameter loop: Event loop
@@ -111,16 +115,31 @@ public final class Skelton {
         Loop.defaultLoop.run()
     }
     
+    public func closeClientsConnected() throws {
+        for client in socketsConnected {
+            try client.close()
+        }
+    }
+    
+    /**
+     Close server handle
+     */
+    public func close() throws {
+        try server.close()
+    }
+    
     private func onConnection(_ queue: PipeSocket?) {
         // TODO need to fix more ARC friendly
         let client = HTTPStream()
-        self.clientsConnected += 1
+        socketsConnected.append(client)
         
         let unmanaged = Unmanaged.passRetained(client)
         
-        client.onClose { [unowned self] in
-            self.clientsConnected -= 1
+        client.onClose { [unowned self, unowned client] in
             unmanaged.release()
+            if let index = self.socketsConnected.index(of: client) {
+                self.socketsConnected.remove(at: index)
+            }
         }
         
         do {
@@ -146,7 +165,7 @@ public final class Skelton {
         
         let parser = RequestParser()
         
-        client.receive { [unowned self] getData in
+        client.receive { [unowned self, unowned client] getData in
             do {
                 let data = try getData()
                 if let request = try parser.parse(data) {
@@ -196,12 +215,5 @@ public final class Skelton {
         
         roundRobinCounter = (roundRobinCounter + 1) % Cluster.workers.count
     }
-    
-    /**
-     Close server handle
-     */
-    public func close() throws {
-        try server.close()
-    }
 }