From 413cafeeb6d2d651482d41e9231650989225e97a Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 7 Oct 2019 17:26:03 +0300 Subject: [PATCH] reworked hub --- hub/connection.go | 87 ++++++++------- hub/connection_hub.go | 109 ++++++++++--------- hub/connection_hub_request.go | 8 -- hub/{closeable_channel.go => ctx_channel.go} | 21 ++-- hub/hub.go | 64 ++++++----- hub/init.go | 30 +++++ hub/registry.go | 18 ++- 7 files changed, 189 insertions(+), 148 deletions(-) delete mode 100644 hub/connection_hub_request.go rename hub/{closeable_channel.go => ctx_channel.go} (62%) create mode 100644 hub/init.go diff --git a/hub/connection.go b/hub/connection.go index 8879f40..17ebbfc 100644 --- a/hub/connection.go +++ b/hub/connection.go @@ -10,34 +10,17 @@ import ( "github.com/9seconds/mtg/protocol" ) -type connectionID int - type connection struct { - conn conntypes.PacketReadWriteCloser - mutex sync.RWMutex - id connectionID - hub *connectionHub - pending uint - closing bool + conn conntypes.PacketReadWriteCloser + mutex sync.RWMutex + shutdownOnce sync.Once + hub *connectionHub + id int + pending uint + done chan struct{} } -func (c *connection) Write(packet conntypes.Packet) error { - c.mutex.Lock() - defer c.mutex.Unlock() - - err := c.conn.Write(packet) - if err != nil { - // if we tried to write into a socket and it was broken, it is - // a time to reconsider the prescence of this socket at all. - // - // probably we need to remove it completely because it seems - // that connection is broken. - c.pending = 0 - } - return err -} - -func (c *connection) Read() (conntypes.Packet, error) { +func (c *connection) read() (conntypes.Packet, error) { packet, err := c.conn.Read() c.mutex.Lock() @@ -51,39 +34,59 @@ func (c *connection) Read() (conntypes.Packet, error) { return packet, err } -func (c *connection) Stats() (bool, uint) { +func (c *connection) write(packet conntypes.Packet) error { + err := c.conn.Write(packet) + if err != nil { + // if we tried to write into a socket and it was broken, it is + // a time to reconsider the prescence of this socket at all. + // + // probably we need to remove it completely because it seems + // that connection is broken. + c.mutex.Lock() + c.pending = 0 + c.mutex.Unlock() + } + return err +} + +func (c *connection) shutdown() { + c.shutdownOnce.Do(func() { + close(c.done) + c.hub.channelBrokenSockets <- c.id + }) +} + +func (c *connection) closed() bool { + select { + case <-c.done: + return true + default: + return false + } +} + +func (c *connection) idle() bool { c.mutex.RLock() defer c.mutex.RUnlock() - return c.closing, c.pending -} - -func (c *connection) Close() error { - c.mutex.Lock() - defer c.mutex.Unlock() - - c.closing = true - return c.conn.Close() + return c.pending == 0 } func (c *connection) run() { for { - packet, err := c.conn.Read() + packet, err := c.read() if err != nil { - c.Close() - c.hub.brokenSocketsChan <- c.id - c.hub = nil + c.shutdown() return } - // TODO if channel, ok := Registry.getChannel(conntypes.ConnID{}); ok { go channel.write(packet) // nolint: errcheck } } } -func newConnection(hub *connectionHub, req *protocol.TelegramRequest) (*connection, error) { +func newConnection(req *protocol.TelegramRequest, hub *connectionHub) (*connection, error) { conn, err := mtproto.TelegramProtocol(req) if err != nil { return nil, fmt.Errorf("cannot create a new connection: %w", err) @@ -92,7 +95,7 @@ func newConnection(hub *connectionHub, req *protocol.TelegramRequest) (*connecti rv := &connection{ conn: conn, hub: hub, - id: connectionID(rand.Int()), + id: rand.Int(), } go rv.run() diff --git a/hub/connection_hub.go b/hub/connection_hub.go index 44bda44..e4aa037 100644 --- a/hub/connection_hub.go +++ b/hub/connection_hub.go @@ -1,84 +1,89 @@ package hub -import "time" +import ( + "time" + + "github.com/9seconds/mtg/protocol" +) const hubGCEvery = time.Minute -type connectionHub struct { - sockets map[connectionID]*connection - - brokenSocketsChan chan connectionID - connectionRequestsChan chan *connectionHubRequest - returnConnectionsChan chan *connection +type connectionHubRequest struct { + request *protocol.TelegramRequest + response chan<- *connection } -func (h *connectionHub) run() { - gcTicker := time.NewTicker(hubGCEvery) - defer gcTicker.Stop() +type connectionHub struct { + sockets map[int]*connection + + channelBrokenSockets chan int + channelConnectionRequests chan *connectionHubRequest + channelReturnConnections chan *connection +} + +func (c *connectionHub) run() { + ticker := time.NewTicker(hubGCEvery) + defer ticker.Stop() for { select { - case <-gcTicker.C: - h.runGC() - case id := <-h.brokenSocketsChan: - h.runBrokenConnection(id) - case request := <-h.connectionRequestsChan: - h.runConnectionRequest(request) - case conn := <-h.returnConnectionsChan: - h.runReturnConnection(conn) + case <-ticker.C: + c.runGC() + case request := <-c.channelConnectionRequests: + c.runConnectionRequest(request) + case id := <-c.channelBrokenSockets: + c.runBrokenSocket(id) + case conn := <-c.channelReturnConnections: + c.runReturnConnection(conn) } } } -func (h *connectionHub) runBrokenConnection(id connectionID) { - delete(h.sockets, id) -} - -func (h *connectionHub) runGC() { - for key, conn := range h.sockets { - closing, pending := conn.Stats() +func (c *connectionHub) runGC() { + for key, conn := range c.sockets { switch { - case closing: - delete(h.sockets, key) - case pending == 0: - conn.Close() - delete(h.sockets, key) + case conn.closed(): + delete(c.sockets, key) + case conn.idle(): + conn.shutdown() + delete(c.sockets, key) return } - } } -func (h *connectionHub) runConnectionRequest(req *connectionHubRequest) { - for key, conn := range h.sockets { - closing, _ := conn.Stats() - delete(h.sockets, key) - - if !closing { - req.responseChan <- conn +func (c *connectionHub) runConnectionRequest(req *connectionHubRequest) { + for key, conn := range c.sockets { + delete(c.sockets, key) + if !conn.closed() { + req.response <- conn + close(req.response) return } } - newConn, err := newConnection(h, req.req) - if err != nil { - close(req.responseChan) - return + if conn, err := newConnection(req.request, c); err == nil { + req.response <- conn } - - req.responseChan <- newConn + close(req.response) } -func (h *connectionHub) runReturnConnection(conn *connection) { - h.sockets[conn.id] = conn +func (c *connectionHub) runBrokenSocket(id int) { + delete(c.sockets, id) +} + +func (c *connectionHub) runReturnConnection(conn *connection) { + c.sockets[conn.id] = conn } func newConnectionHub() *connectionHub { - return &connectionHub{ - sockets: map[connectionID]*connection{}, - - brokenSocketsChan: make(chan connectionID, 1), - connectionRequestsChan: make(chan *connectionHubRequest), - returnConnectionsChan: make(chan *connection, 1), + rv := &connectionHub{ + sockets: map[int]*connection{}, + channelBrokenSockets: make(chan int, 1), + channelConnectionRequests: make(chan *connectionHubRequest), + channelReturnConnections: make(chan *connection, 1), } + go rv.run() + + return rv } diff --git a/hub/connection_hub_request.go b/hub/connection_hub_request.go deleted file mode 100644 index eef2c7c..0000000 --- a/hub/connection_hub_request.go +++ /dev/null @@ -1,8 +0,0 @@ -package hub - -import "github.com/9seconds/mtg/protocol" - -type connectionHubRequest struct { - req *protocol.TelegramRequest - responseChan chan<- *connection -} diff --git a/hub/closeable_channel.go b/hub/ctx_channel.go similarity index 62% rename from hub/closeable_channel.go rename to hub/ctx_channel.go index 5042ef7..9394ba3 100644 --- a/hub/closeable_channel.go +++ b/hub/ctx_channel.go @@ -12,46 +12,47 @@ const closeableChannelReadTimeout = 2 * time.Minute type ChannelReadCloser interface { Read() (conntypes.Packet, error) - Close() + Close() error } -type closeableChannel struct { +type ctxChannel struct { channel chan conntypes.Packet ctx context.Context cancel context.CancelFunc } -func (c *closeableChannel) Read() (conntypes.Packet, error) { +func (c *ctxChannel) Read() (conntypes.Packet, error) { timer := time.NewTimer(closeableChannelReadTimeout) defer timer.Stop() select { case <-timer.C: - return nil, errors.New("timeout") + return nil, ErrTimeout case <-c.ctx.Done(): - return nil, errors.New("channel was closed") + return nil, ErrClosed case packet := <-c.channel: return packet, nil } } -func (c *closeableChannel) write(packet conntypes.Packet) error { +func (c *ctxChannel) write(packet conntypes.Packet) error { select { case <-c.ctx.Done(): - return errors.New("channel was closed") + return ErrClosed case c.channel <- packet: return nil } } -func (c *closeableChannel) Close() { +func (c *ctxChannel) Close() error { c.cancel() c.channel = nil + return nil } -func newCloseableChannel(ctx context.Context) *closeableChannel { +func newCtxChannel(ctx context.Context) *ctxChannel { ctx, cancel := context.WithCancel(ctx) - return &closeableChannel{ + return &ctxChannel{ channel: make(chan conntypes.Packet), ctx: ctx, cancel: cancel, diff --git a/hub/hub.go b/hub/hub.go index ef160d0..1e6cd5d 100644 --- a/hub/hub.go +++ b/hub/hub.go @@ -1,48 +1,60 @@ package hub import ( - "errors" + "encoding/binary" + "fmt" + "strings" "sync" "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/protocol" ) -type Concentrator struct { - hubs sync.Map +type hub struct { + subs map[string]*connectionHub + mutex sync.RWMutex } -func (c *Concentrator) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error { - hub := c.getHub(req) - connectionChan := make(chan *connection) - hub.connectionRequestsChan <- &connectionHubRequest{ - req: req, - responseChan: connectionChan, +func (h *hub) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error { + sub := h.getHub(req) + connections := make(chan *connection) + sub.channelConnectionRequests <- &connectionHubRequest{ + request: req, + response: connections, } - conn, ok := <-connectionChan + conn, ok := <-connections if !ok { - return errors.New("cannot establish connection to telegram") + return ErrCannotCreateConnection } + + if err := conn.write(packet); err != nil { + return fmt.Errorf("cannot send packet: %w", err) + } + return nil } -func (c *Concentrator) getHub(req *protocol.TelegramRequest) *connectionHub { - dcMapRaw, ok := c.hubs.Load(req.ClientProtocol.DC()) - if !ok { - dcMapRaw, _ = c.hubs.LoadOrStore(req.ClientProtocol.DC(), &sync.Map{}) - } - dcMap := dcMapRaw.(*sync.Map) +func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub { + keyBuilder := strings.Builder{} + binary.Write(&keyBuilder, binary.LittleEndian, int16(req.ClientProtocol.DC())) + keyBuilder.WriteRune('_') + binary.Write(&keyBuilder, binary.LittleEndian, uint8(req.ClientProtocol.ConnectionProtocol())) + key := keyBuilder.String() + + h.mutex.RLock() + rv, ok := h.subs[key] + h.mutex.RUnlock() - loaded := true - hubRaw, ok := dcMap.Load(req.ClientProtocol.ConnectionProtocol()) if !ok { - hubRaw, loaded = dcMap.LoadOrStore(req.ClientProtocol.ConnectionProtocol(), - newConnectionHub()) - } - hub := hubRaw.(*connectionHub) - if !loaded { - go hub.run() + h.mutex.Lock() + defer h.mutex.Unlock() + + rv, ok = h.subs[key] + if !ok { + rv = newConnectionHub() + h.subs[key] = rv + } } - return hub + return rv } diff --git a/hub/init.go b/hub/init.go new file mode 100644 index 0000000..38d1a82 --- /dev/null +++ b/hub/init.go @@ -0,0 +1,30 @@ +package hub + +import ( + "context" + "errors" + "sync" +) + +var ( + Registry *registry + Hub *hub + + ErrTimeout = errors.New("timeout") + ErrClosed = errors.New("channel was closed") + ErrCannotCreateConnection = errors.New("cannot create connection") + + initOnce sync.Once +) + +func Init(ctx context.Context) { + initOnce.Do(func() { + Registry = ®istry{ + conns: map[string]*ctxChannel{}, + ctx: ctx, + } + Hub = &hub{ + subs: map[string]*connectionHub{}, + } + }) +} diff --git a/hub/registry.go b/hub/registry.go index 0891e45..4beebe8 100644 --- a/hub/registry.go +++ b/hub/registry.go @@ -7,16 +7,14 @@ import ( "github.com/9seconds/mtg/conntypes" ) -var Registry *RegistryStruct - -type RegistryStruct struct { - conns map[string]*closeableChannel +type registry struct { + conns map[string]*ctxChannel ctx context.Context mutex sync.RWMutex } -func (r *RegistryStruct) Register(id conntypes.ConnID) ChannelReadCloser { - channel := newCloseableChannel(r.ctx) +func (r *registry) Register(id conntypes.ConnID) ChannelReadCloser { + channel := newCtxChannel(r.ctx) r.mutex.Lock() r.conns[string(id[:])] = channel @@ -25,7 +23,7 @@ func (r *RegistryStruct) Register(id conntypes.ConnID) ChannelReadCloser { return channel } -func (r *RegistryStruct) Unregister(id conntypes.ConnID) { +func (r *registry) Unregister(id conntypes.ConnID) { r.mutex.Lock() defer r.mutex.Unlock() @@ -35,7 +33,7 @@ func (r *RegistryStruct) Unregister(id conntypes.ConnID) { } } -func (r *RegistryStruct) getChannel(id conntypes.ConnID) (*closeableChannel, bool) { +func (r *registry) getChannel(id conntypes.ConnID) (*ctxChannel, bool) { r.mutex.RLock() defer r.mutex.RUnlock() @@ -46,8 +44,8 @@ func (r *RegistryStruct) getChannel(id conntypes.ConnID) (*closeableChannel, boo } func InitRegistry(ctx context.Context) { - Registry = &RegistryStruct{ + Registry = ®istry{ ctx: ctx, - conns: map[string]*closeableChannel{}, + conns: map[string]*ctxChannel{}, } }