delete unused connections from the multiplexer

mock_multiplexer_test.go

@@ -46,3 +46,15 @@ func (m *MockMultiplexer) AddConn(arg0 net.PacketConn, arg1 int) (packetHandlerM
 func (mr *MockMultiplexerMockRecorder) AddConn(arg0, arg1 interface{}) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddConn", reflect.TypeOf((*MockMultiplexer)(nil).AddConn), arg0, arg1)
 }
+
+// RemoveConn mocks base method
+func (m *MockMultiplexer) RemoveConn(arg0 net.PacketConn) error {
+	ret := m.ctrl.Call(m, "RemoveConn", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RemoveConn indicates an expected call of RemoveConn
+func (mr *MockMultiplexerMockRecorder) RemoveConn(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveConn", reflect.TypeOf((*MockMultiplexer)(nil).RemoveConn), arg0)
+}

multiplexer.go

@@ -15,6 +15,7 @@ var (
 
 type multiplexer interface {
 	AddConn(net.PacketConn, int) (packetHandlerManager, error)
+	RemoveConn(net.PacketConn) error
 }
 
 type connManager struct {
@@ -61,3 +62,15 @@ func (m *connMultiplexer) AddConn(c net.PacketConn, connIDLen int) (packetHandle
 	}
 	return p.manager, nil
 }
+
+func (m *connMultiplexer) RemoveConn(c net.PacketConn) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if _, ok := m.conns[c]; !ok {
+		return fmt.Errorf("cannote remove connection, connection is unknown")
+	}
+
+	delete(m.conns, c)
+	return nil
+}

packet_handler_map.go

@@ -139,7 +139,7 @@ func (h *packetHandlerMap) close(e error) error {
 	}
 	h.mutex.Unlock()
 	wg.Wait()
-	return nil
+	return getMultiplexer().RemoveConn(h.conn)
 }
 
 func (h *packetHandlerMap) listen() {

packet_handler_map_test.go

@@ -45,6 +45,16 @@ var _ = Describe("Packet Handler Map", func() {
 	})
 
 	It("closes", func() {
+		getMultiplexer() // make the sync.Once execute
+		// replace the clientMuxer. getClientMultiplexer will now return the MockMultiplexer
+		mockMultiplexer := NewMockMultiplexer(mockCtrl)
+		origMultiplexer := connMuxer
+		connMuxer = mockMultiplexer
+
+		defer func() {
+			connMuxer = origMultiplexer
+		}()
+
 		testErr := errors.New("test error	")
 		sess1 := NewMockPacketHandler(mockCtrl)
 		sess1.EXPECT().destroy(testErr)
@@ -52,6 +62,7 @@ var _ = Describe("Packet Handler Map", func() {
 		sess2.EXPECT().destroy(testErr)
 		handler.Add(protocol.ConnectionID{1, 1, 1, 1}, sess1)
 		handler.Add(protocol.ConnectionID{2, 2, 2, 2}, sess2)
+		mockMultiplexer.EXPECT().RemoveConn(gomock.Any())
 		handler.close(testErr)
 	})