ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 11:44:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

reworked hub

Browse Source
This commit is contained in:
9seconds
2019-10-07 17:26:03 +03:00
parent c9743b5675
commit 413cafeeb6
7 changed files with 189 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hub/connection.go
+45 -42
View File
@@ -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()
hub/connection_hub.go
+57 -52
View File
@@ -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
}
hub/connection_hub_request.go
-8
View File
@@ -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
+11 -10
View File
@@ -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
+38 -26
View File
@@ -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
}
hub/init.go
+30
View File
@@ -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
+8 -10
View File
@@ -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{},
}
}