reworked hub

hub/connection.go

@@ -10,21 +10,31 @@ 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 {
+func (c *connection) read() (conntypes.Packet, error) {
+	packet, err := c.conn.Read()
+
 	c.mutex.Lock()
-	defer c.mutex.Unlock()
+	if err != nil {
+		c.pending--
+	} else {
+		c.pending = 0
+	}
+	c.mutex.Unlock()
+
+	return packet, err
+}
 
+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
@@ -32,58 +42,51 @@ func (c *connection) Write(packet conntypes.Packet) error {
 		//
 		// 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) Read() (conntypes.Packet, error) {
-	packet, err := c.conn.Read()
+func (c *connection) shutdown() {
+	c.shutdownOnce.Do(func() {
+		close(c.done)
+			c.hub.channelBrokenSockets <- c.id
+	})
+}
 
-	c.mutex.Lock()
-	if err != nil {
-		c.pending--
-	} else {
-		c.pending = 0
+func (c *connection) closed() bool {
+	select {
+	case <-c.done:
+		return true
+	default:
+		return false
 	}
-	c.mutex.Unlock()
-
-	return packet, err
 }
 
-func (c *connection) Stats() (bool, uint) {
+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()
 

hub/connection_hub.go

@@ -1,84 +1,89 @@
 package hub
 
-import "time"
+import (
+	"time"
+
+	"github.com/9seconds/mtg/protocol"
+)
 
 const hubGCEvery = time.Minute
 
+type connectionHubRequest struct {
+	request  *protocol.TelegramRequest
+	response chan<- *connection
+}
+
 type connectionHub struct {
-	sockets map[connectionID]*connection
+	sockets map[int]*connection
 
-	brokenSocketsChan      chan connectionID
-	connectionRequestsChan chan *connectionHubRequest
-	returnConnectionsChan  chan *connection
+	channelBrokenSockets      chan int
+	channelConnectionRequests chan *connectionHubRequest
+	channelReturnConnections  chan *connection
 }
 
-func (h *connectionHub) run() {
-	gcTicker := time.NewTicker(hubGCEvery)
-	defer gcTicker.Stop()
+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
 	}
+	close(req.response)
+}
 
-	req.responseChan <- newConn
+func (c *connectionHub) runBrokenSocket(id int) {
+	delete(c.sockets, id)
 }
 
-func (h *connectionHub) runReturnConnection(conn *connection) {
-	h.sockets[conn.id] = conn
+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
 }

hub/connection_hub_request.go

@@ -1,8 +0,0 @@
-package hub
-
-import "github.com/9seconds/mtg/protocol"
-
-type connectionHubRequest struct {
-	req          *protocol.TelegramRequest
-	responseChan chan<- *connection
-}

hub/closeable_channel.go → 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,

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
 	}
-}
 
-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{})
+	if err := conn.write(packet); err != nil {
+		return fmt.Errorf("cannot send packet: %w", err)
 	}
-	dcMap := dcMapRaw.(*sync.Map)
+	return nil
+}
+
+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
 }

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 = &registry{
+			conns: map[string]*ctxChannel{},
+			ctx:   ctx,
+		}
+		Hub = &hub{
+			subs: map[string]*connectionHub{},
+		}
+	})
+}

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 = &registry{
 		ctx:   ctx,
-		conns: map[string]*closeableChannel{},
+		conns: map[string]*ctxChannel{},
 	}
 }