wip

hub/closeable_channel.go

@@ -0,0 +1,59 @@
+package hub
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+const closeableChannelReadTimeout = 2 * time.Minute
+
+type ChannelReadCloser interface {
+	Read() (conntypes.Packet, error)
+	Close()
+}
+
+type closeableChannel struct {
+	channel chan conntypes.Packet
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+func (c *closeableChannel) Read() (conntypes.Packet, error) {
+	timer := time.NewTimer(closeableChannelReadTimeout)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return nil, errors.New("timeout")
+	case <-c.ctx.Done():
+		return nil, errors.New("channel was closed")
+	case packet := <-c.channel:
+		return packet, nil
+	}
+}
+
+func (c *closeableChannel) write(packet conntypes.Packet) error {
+	select {
+	case <-c.ctx.Done():
+		return errors.New("channel was closed")
+	case c.channel <- packet:
+		return nil
+	}
+}
+
+func (c *closeableChannel) Close() {
+	c.cancel()
+	c.channel = nil
+}
+
+func newCloseableChannel(ctx context.Context) *closeableChannel {
+	ctx, cancel := context.WithCancel(ctx)
+	return &closeableChannel{
+		channel: make(chan conntypes.Packet),
+		ctx:     ctx,
+		cancel:  cancel,
+	}
+}

hub/connection.go

@@ -0,0 +1,100 @@
+package hub
+
+import (
+	"fmt"
+	"math/rand"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/mtproto"
+	"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
+}
+
+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) {
+	packet, err := c.conn.Read()
+
+	c.mutex.Lock()
+	if err != nil {
+		c.pending--
+	} else {
+		c.pending = 0
+	}
+	c.mutex.Unlock()
+
+	return packet, err
+}
+
+func (c *connection) Stats() (bool, uint) {
+	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()
+}
+
+func (c *connection) run() {
+	for {
+		packet, err := c.conn.Read()
+		if err != nil {
+			c.Close()
+			c.hub.brokenSocketsChan <- c.id
+			c.hub = nil
+			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) {
+	conn, err := mtproto.TelegramProtocol(req)
+	if err != nil {
+		return nil, fmt.Errorf("cannot create a new connection: %w", err)
+	}
+
+	rv := &connection{
+		conn: conn,
+		hub:  hub,
+		id:   connectionID(rand.Int()),
+	}
+	go rv.run()
+
+	return rv, nil
+}

hub/connection_hub.go

@@ -0,0 +1,84 @@
+package hub
+
+import "time"
+
+const hubGCEvery = time.Minute
+
+type connectionHub struct {
+	sockets map[connectionID]*connection
+
+	brokenSocketsChan      chan connectionID
+	connectionRequestsChan chan *connectionHubRequest
+	returnConnectionsChan  chan *connection
+}
+
+func (h *connectionHub) run() {
+	gcTicker := time.NewTicker(hubGCEvery)
+	defer gcTicker.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)
+		}
+	}
+}
+
+func (h *connectionHub) runBrokenConnection(id connectionID) {
+	delete(h.sockets, id)
+}
+
+func (h *connectionHub) runGC() {
+	for key, conn := range h.sockets {
+		closing, pending := conn.Stats()
+		switch {
+		case closing:
+			delete(h.sockets, key)
+		case pending == 0:
+			conn.Close()
+			delete(h.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
+			return
+		}
+	}
+
+	newConn, err := newConnection(h, req.req)
+	if err != nil {
+		close(req.responseChan)
+		return
+	}
+
+	req.responseChan <- newConn
+}
+
+func (h *connectionHub) runReturnConnection(conn *connection) {
+	h.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),
+	}
+}

hub/connection_hub_request.go

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

hub/hub.go

@@ -0,0 +1,48 @@
+package hub
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/protocol"
+)
+
+type Concentrator struct {
+	hubs sync.Map
+}
+
+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,
+	}
+
+	conn, ok := <-connectionChan
+	if !ok {
+		return errors.New("cannot establish connection to telegram")
+	}
+}
+
+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)
+
+	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()
+	}
+
+	return hub
+}

hub/registry.go

@@ -0,0 +1,53 @@
+package hub
+
+import (
+	"context"
+	"sync"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+var Registry *RegistryStruct
+
+type RegistryStruct struct {
+	conns map[string]*closeableChannel
+	ctx   context.Context
+	mutex sync.RWMutex
+}
+
+func (r *RegistryStruct) Register(id conntypes.ConnID) ChannelReadCloser {
+	channel := newCloseableChannel(r.ctx)
+
+	r.mutex.Lock()
+	r.conns[string(id[:])] = channel
+	r.mutex.Unlock()
+
+	return channel
+}
+
+func (r *RegistryStruct) Unregister(id conntypes.ConnID) {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if channel, ok := r.conns[string(id[:])]; ok {
+		channel.Close()
+		delete(r.conns, string(id[:]))
+	}
+}
+
+func (r *RegistryStruct) getChannel(id conntypes.ConnID) (*closeableChannel, bool) {
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+
+	if value, ok := r.conns[string(id[:])]; ok {
+		return value, true
+	}
+	return nil, false
+}
+
+func InitRegistry(ctx context.Context) {
+	Registry = &RegistryStruct{
+		ctx:   ctx,
+		conns: map[string]*closeableChannel{},
+	}
+}