Add unregistering on telegram conn close

This commit is contained in:
9seconds
2019-10-10 16:42:29 +03:00
parent 61450ea3db
commit 2eba78b0db
2 changed files with 10 additions and 1 deletions
+3
View File
@@ -54,6 +54,9 @@ func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub {
rv, ok = h.subs[key] rv, ok = h.subs[key]
if !ok { if !ok {
h.logger.Debugw("Create new connection hub",
"dc", req.ClientProtocol.DC(),
"protocol", req.ClientProtocol.ConnectionProtocol())
rv = newConnectionHub(h.logger.With( rv = newConnectionHub(h.logger.With(
"dc", req.ClientProtocol.DC(), "dc", req.ClientProtocol.DC(),
"protocol", req.ClientProtocol.ConnectionProtocol(), "protocol", req.ClientProtocol.ConnectionProtocol(),
+7 -1
View File
@@ -5,6 +5,7 @@ import (
"encoding/binary" "encoding/binary"
"fmt" "fmt"
"net" "net"
"sync"
"github.com/9seconds/mtg/config" "github.com/9seconds/mtg/config"
"github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/conntypes"
@@ -19,6 +20,7 @@ type wrapperProxy struct {
clientIPPort []byte clientIPPort []byte
ourIPPort []byte ourIPPort []byte
channelRead hub.ChannelReadCloser channelRead hub.ChannelReadCloser
closeOnce sync.Once
} }
func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error { func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error {
@@ -61,7 +63,11 @@ func (w *wrapperProxy) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, e
} }
func (w *wrapperProxy) Close() error { func (w *wrapperProxy) Close() error {
return w.channelRead.Close() w.closeOnce.Do(func() {
w.channelRead.Close()
hub.Registry.Unregister(w.request.ConnID)
})
return nil
} }
func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteCloser { func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteCloser {