mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-09-01 05:04:03 +03:00
Correct multiplexing
This commit is contained in:
9seconds
+117
-74
@@ -5,6 +5,8 @@ import (
|
||||
"math/rand"
|
||||
"sync"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/mtproto"
|
||||
"github.com/9seconds/mtg/mtproto/rpc"
|
||||
@@ -12,108 +14,149 @@ import (
|
||||
)
|
||||
|
||||
type connection struct {
|
||||
conn conntypes.PacketReadWriteCloser
|
||||
mutex sync.RWMutex
|
||||
shutdownOnce sync.Once
|
||||
hub *connectionHub
|
||||
id int
|
||||
pending uint
|
||||
done chan struct{}
|
||||
}
|
||||
conn conntypes.PacketReadWriteCloser
|
||||
proxyConns map[string]*ProxyConn
|
||||
closeOnce sync.Once
|
||||
proxyConnsMutex sync.RWMutex
|
||||
id int
|
||||
logger *zap.SugaredLogger
|
||||
|
||||
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) 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() {
|
||||
c.conn.Close()
|
||||
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.pending == 0
|
||||
channelDone chan struct{}
|
||||
channelWrite chan conntypes.Packet
|
||||
channelRead chan *rpc.ProxyResponse
|
||||
channelConnAttach chan *ProxyConn
|
||||
channelConnDetach chan conntypes.ConnID
|
||||
}
|
||||
|
||||
func (c *connection) run() {
|
||||
logger := c.hub.logger.Named("connection").With("id", c.id)
|
||||
defer c.Close()
|
||||
|
||||
for {
|
||||
packet, err := c.read()
|
||||
select {
|
||||
case <-c.channelDone:
|
||||
for _, v := range c.proxyConns {
|
||||
v.Close()
|
||||
}
|
||||
|
||||
return
|
||||
case resp := <-c.channelRead:
|
||||
if channel, ok := c.proxyConns[string(resp.ConnID[:])]; ok {
|
||||
if resp.Type == rpc.ProxyResponseTypeCloseExt {
|
||||
channel.Close()
|
||||
} else {
|
||||
channel.put(resp)
|
||||
}
|
||||
}
|
||||
case packet := <-c.channelWrite:
|
||||
if err := c.conn.Write(packet); err != nil {
|
||||
c.logger.Debugw("Cannot write packet", "error", err)
|
||||
c.Close()
|
||||
}
|
||||
case conn := <-c.channelConnAttach:
|
||||
c.proxyConnsMutex.Lock()
|
||||
c.proxyConns[string(conn.req.ConnID[:])] = conn
|
||||
c.proxyConnsMutex.Unlock()
|
||||
conn.channelWrite = c.channelWrite
|
||||
case connID := <-c.channelConnDetach:
|
||||
if conn, ok := c.proxyConns[string(connID[:])]; ok {
|
||||
c.proxyConnsMutex.Lock()
|
||||
delete(c.proxyConns, string(connID[:]))
|
||||
c.proxyConnsMutex.Unlock()
|
||||
conn.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *connection) readLoop() {
|
||||
for {
|
||||
packet, err := c.conn.Read()
|
||||
if err != nil {
|
||||
c.shutdown()
|
||||
c.logger.Debugw("Cannot read packet", "error", err)
|
||||
c.Close()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
response, err := rpc.ParseProxyResponse(packet)
|
||||
if err != nil {
|
||||
logger.Debugw("Failed response", "error", err)
|
||||
c.logger.Debugw("Failed response", "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
if response.Type == rpc.ProxyResponseTypeCloseExt {
|
||||
logger.Debugw("Proxy has closed connection")
|
||||
select {
|
||||
case <-c.channelDone:
|
||||
return
|
||||
}
|
||||
|
||||
if channel, ok := Registry.getChannel(response.ConnID); ok {
|
||||
go channel.sendBack(response) // nolint: errcheck
|
||||
case c.channelRead <- response:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func newConnection(req *protocol.TelegramRequest, hub *connectionHub) (*connection, error) {
|
||||
func (c *connection) Close() {
|
||||
c.closeOnce.Do(func() {
|
||||
c.logger.Debugw("Closing connection")
|
||||
|
||||
close(c.channelDone)
|
||||
c.conn.Close()
|
||||
})
|
||||
}
|
||||
|
||||
func (c *connection) Done() bool {
|
||||
select {
|
||||
case <-c.channelDone:
|
||||
return true
|
||||
default:
|
||||
return c.Len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
func (c *connection) Len() int {
|
||||
c.proxyConnsMutex.RLock()
|
||||
defer c.proxyConnsMutex.RUnlock()
|
||||
|
||||
return len(c.proxyConns)
|
||||
}
|
||||
|
||||
func (c *connection) Attach(conn *ProxyConn) error {
|
||||
select {
|
||||
case <-c.channelDone:
|
||||
return ErrClosed
|
||||
case c.channelConnAttach <- conn:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *connection) Detach(connID conntypes.ConnID) {
|
||||
select {
|
||||
case <-c.channelDone:
|
||||
case c.channelConnDetach <- connID:
|
||||
}
|
||||
}
|
||||
|
||||
func newConnection(req *protocol.TelegramRequest) (*connection, error) {
|
||||
conn, err := mtproto.TelegramProtocol(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot create a new connection: %w", err)
|
||||
}
|
||||
|
||||
id := rand.Int() // nolint: gosec
|
||||
rv := &connection{
|
||||
conn: conn,
|
||||
hub: hub,
|
||||
id: rand.Int(), // nolint: gosec
|
||||
done: make(chan struct{}),
|
||||
id: id,
|
||||
logger: zap.S().Named("hub-connection").With("id", id,
|
||||
"dc", req.ClientProtocol.DC(),
|
||||
"protocol", req.ClientProtocol.ConnectionProtocol()),
|
||||
proxyConns: make(map[string]*ProxyConn),
|
||||
|
||||
channelRead: make(chan *rpc.ProxyResponse, 1),
|
||||
channelDone: make(chan struct{}),
|
||||
channelWrite: make(chan conntypes.Packet),
|
||||
channelConnAttach: make(chan *ProxyConn),
|
||||
channelConnDetach: make(chan conntypes.ConnID),
|
||||
}
|
||||
|
||||
go rv.readLoop()
|
||||
|
||||
go rv.run()
|
||||
|
||||
return rv, nil
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/protocol"
|
||||
)
|
||||
|
||||
const hubGCEvery = time.Minute
|
||||
|
||||
type connectionHubRequest struct {
|
||||
request *protocol.TelegramRequest
|
||||
response chan<- *connection
|
||||
}
|
||||
|
||||
type connectionHub struct {
|
||||
sockets map[int]*connection
|
||||
logger *zap.SugaredLogger
|
||||
|
||||
channelBrokenSockets chan int
|
||||
channelConnectionRequests chan *connectionHubRequest
|
||||
channelReturnConnections chan *connection
|
||||
}
|
||||
|
||||
func (c *connectionHub) run() {
|
||||
ticker := time.NewTicker(hubGCEvery)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
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 (c *connectionHub) runGC() {
|
||||
logger := c.logger.Named("gc")
|
||||
|
||||
for key, conn := range c.sockets {
|
||||
switch {
|
||||
case conn.closed():
|
||||
logger.Debugw("Delete closed socket", "key", key)
|
||||
delete(c.sockets, key)
|
||||
case conn.idle():
|
||||
logger.Debugw("Delete idle socket", "key", key)
|
||||
conn.shutdown()
|
||||
delete(c.sockets, key)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *connectionHub) runConnectionRequest(req *connectionHubRequest) {
|
||||
logger := c.logger.Named("request").With("connection-id", req.request.ConnID)
|
||||
|
||||
for key, conn := range c.sockets {
|
||||
delete(c.sockets, key)
|
||||
|
||||
if !conn.closed() {
|
||||
logger.Debugw("Choose connection",
|
||||
"id", conn.id,
|
||||
"remote_addr", conn.conn.RemoteAddr())
|
||||
req.response <- conn
|
||||
close(req.response)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if conn, err := newConnection(req.request, c); err == nil {
|
||||
logger.Debugw("New connection",
|
||||
"id", conn.id,
|
||||
"remote_addr", conn.conn.RemoteAddr())
|
||||
req.response <- conn
|
||||
}
|
||||
|
||||
close(req.response)
|
||||
}
|
||||
|
||||
func (c *connectionHub) runBrokenSocket(id int) {
|
||||
c.logger.Named("broken-socket").Debugw("Delete broken socket", "id", id)
|
||||
delete(c.sockets, id)
|
||||
}
|
||||
|
||||
func (c *connectionHub) runReturnConnection(conn *connection) {
|
||||
c.logger.Named("return-connection").Debugw("Return connection",
|
||||
"id", conn.id,
|
||||
"remote_addr", conn.conn.RemoteAddr())
|
||||
|
||||
c.sockets[conn.id] = conn
|
||||
}
|
||||
|
||||
func newConnectionHub(logger *zap.SugaredLogger) *connectionHub {
|
||||
rv := &connectionHub{
|
||||
logger: logger.Named("connection-hub"),
|
||||
sockets: map[int]*connection{},
|
||||
channelBrokenSockets: make(chan int, 1),
|
||||
channelConnectionRequests: make(chan *connectionHubRequest),
|
||||
channelReturnConnections: make(chan *connection, 1),
|
||||
}
|
||||
go rv.run()
|
||||
|
||||
return rv
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
)
|
||||
|
||||
const connectionListMaxClientsPerConnection = 2
|
||||
|
||||
type connectionList struct {
|
||||
connections []*connection
|
||||
}
|
||||
|
||||
func (c *connectionList) Get(conn *ProxyConn) (*connection, error) {
|
||||
if len(c.connections) > 0 {
|
||||
c.gc()
|
||||
}
|
||||
|
||||
if len(c.connections) > 0 && c.connections[0].Len() < connectionListMaxClientsPerConnection {
|
||||
if err := c.connections[0].Attach(conn); err == nil {
|
||||
return c.connections[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
newConn, err := newConnection(conn.req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot allocate a new connection: %w", err)
|
||||
}
|
||||
|
||||
if err = newConn.Attach(conn); err != nil {
|
||||
newConn.Close()
|
||||
return nil, fmt.Errorf("cannot attach to the newly created connection: %w", err)
|
||||
}
|
||||
|
||||
c.connections = append(c.connections, newConn)
|
||||
lastIndex := len(c.connections) - 1
|
||||
c.connections[0], c.connections[lastIndex] = c.connections[lastIndex], c.connections[0]
|
||||
|
||||
return newConn, nil
|
||||
}
|
||||
|
||||
func (c *connectionList) gc() {
|
||||
prevLen := len(c.connections)
|
||||
|
||||
for i := len(c.connections) - 1; i >= 0; i-- {
|
||||
lastIndex := len(c.connections) - 1
|
||||
|
||||
if c.connections[i].Done() {
|
||||
c.connections[i].Close()
|
||||
|
||||
if len(c.connections)-1 == i {
|
||||
c.connections = c.connections[:lastIndex]
|
||||
} else {
|
||||
c.connections[i], c.connections[lastIndex] = c.connections[lastIndex], c.connections[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if prevLen != len(c.connections) {
|
||||
c.sort()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *connectionList) sort() {
|
||||
if len(c.connections) > 1 {
|
||||
sort.Slice(c.connections, func(i, j int) bool {
|
||||
return c.connections[i].Len() < c.connections[j].Len()
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/mtproto/rpc"
|
||||
)
|
||||
|
||||
const closeableChannelReadTimeout = 2 * time.Minute
|
||||
|
||||
type ChannelReadCloser interface {
|
||||
Read() (*rpc.ProxyResponse, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
type ctxChannel struct {
|
||||
channel chan *rpc.ProxyResponse
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func (c *ctxChannel) Read() (*rpc.ProxyResponse, error) {
|
||||
timer := time.NewTimer(closeableChannelReadTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-timer.C:
|
||||
return nil, ErrTimeout
|
||||
case <-c.ctx.Done():
|
||||
return nil, ErrClosed
|
||||
case packet := <-c.channel:
|
||||
return packet, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *ctxChannel) sendBack(response *rpc.ProxyResponse) error {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return ErrClosed
|
||||
case c.channel <- response:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *ctxChannel) Close() error {
|
||||
c.cancel()
|
||||
c.channel = nil
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func newCtxChannel(ctx context.Context) *ctxChannel {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
|
||||
return &ctxChannel{
|
||||
channel: make(chan *rpc.ProxyResponse),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
+15
-48
@@ -1,73 +1,40 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"strings"
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/protocol"
|
||||
)
|
||||
|
||||
type hub struct {
|
||||
logger *zap.SugaredLogger
|
||||
subs map[string]*connectionHub
|
||||
mutex sync.RWMutex
|
||||
muxes map[int32]*mux
|
||||
mutex sync.RWMutex
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
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 := <-connections
|
||||
if !ok {
|
||||
return ErrCannotCreateConnection
|
||||
}
|
||||
|
||||
if err := conn.write(packet); err != nil {
|
||||
conn.shutdown()
|
||||
return fmt.Errorf("cannot send packet: %w", err)
|
||||
}
|
||||
sub.channelReturnConnections <- conn
|
||||
|
||||
return nil
|
||||
func (h *hub) Register(req *protocol.TelegramRequest) (*ProxyConn, error) {
|
||||
return h.getMux(req).Get(req)
|
||||
}
|
||||
|
||||
func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub {
|
||||
keyBuilder := strings.Builder{}
|
||||
binary.Write(&keyBuilder, binary.LittleEndian, int16(req.ClientProtocol.DC())) // nolint: errcheck
|
||||
keyBuilder.WriteRune('_')
|
||||
binary.Write(&keyBuilder, binary.LittleEndian, uint8(req.ClientProtocol.ConnectionProtocol())) // nolint: errcheck
|
||||
key := keyBuilder.String()
|
||||
func (h *hub) getMux(req *protocol.TelegramRequest) *mux {
|
||||
var key int32 = 32767 + int32(req.ClientProtocol.DC()) + 100000*int32(req.ClientProtocol.ConnectionProtocol())
|
||||
|
||||
h.mutex.RLock()
|
||||
rv, ok := h.subs[key]
|
||||
m, ok := h.muxes[key]
|
||||
h.mutex.RUnlock()
|
||||
|
||||
if !ok {
|
||||
h.mutex.Lock()
|
||||
defer h.mutex.Unlock()
|
||||
m, ok = h.muxes[key]
|
||||
|
||||
rv, ok = h.subs[key]
|
||||
if !ok {
|
||||
h.logger.Debugw("Create new connection hub",
|
||||
"dc", req.ClientProtocol.DC(),
|
||||
"protocol", req.ClientProtocol.ConnectionProtocol())
|
||||
|
||||
rv = newConnectionHub(h.logger.With(
|
||||
"dc", req.ClientProtocol.DC(),
|
||||
"protocol", req.ClientProtocol.ConnectionProtocol(),
|
||||
))
|
||||
h.subs[key] = rv
|
||||
m = newMux(h.ctx)
|
||||
h.muxes[key] = m
|
||||
}
|
||||
|
||||
h.mutex.Unlock()
|
||||
}
|
||||
|
||||
return rv
|
||||
return m
|
||||
}
|
||||
|
||||
+5
-14
@@ -4,30 +4,21 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
var (
|
||||
Registry *registry
|
||||
Hub *hub
|
||||
|
||||
ErrTimeout = errors.New("timeout")
|
||||
ErrClosed = errors.New("channel was closed")
|
||||
ErrCannotCreateConnection = errors.New("cannot create connection")
|
||||
ErrTimeout = errors.New("timeout")
|
||||
ErrClosed = errors.New("context is closed")
|
||||
|
||||
Hub Interface
|
||||
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{},
|
||||
logger: zap.S().Named("hub"),
|
||||
muxes: make(map[int32]*mux),
|
||||
ctx: ctx,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
package hub
|
||||
|
||||
import "github.com/9seconds/mtg/protocol"
|
||||
|
||||
type Interface interface {
|
||||
Register(*protocol.TelegramRequest) (*ProxyConn, error)
|
||||
}
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/protocol"
|
||||
)
|
||||
|
||||
type muxNewRequest struct {
|
||||
req *protocol.TelegramRequest
|
||||
resp chan<- muxNewResponse
|
||||
}
|
||||
|
||||
type muxNewResponse struct {
|
||||
conn *ProxyConn
|
||||
err error
|
||||
}
|
||||
|
||||
type mux struct {
|
||||
connections connectionList
|
||||
clients map[string]*connection
|
||||
ctx context.Context
|
||||
channelClosed chan conntypes.ConnID
|
||||
channelNew chan muxNewRequest
|
||||
}
|
||||
|
||||
func (m *mux) run() {
|
||||
for {
|
||||
select {
|
||||
case <-m.ctx.Done():
|
||||
for _, v := range m.clients {
|
||||
v.Close()
|
||||
}
|
||||
|
||||
return
|
||||
case req := <-m.channelNew:
|
||||
proxyConn := newProxyConn(req.req, m.channelClosed)
|
||||
conn, err := m.connections.Get(proxyConn)
|
||||
|
||||
if err == nil {
|
||||
m.clients[string(req.req.ConnID[:])] = conn
|
||||
}
|
||||
|
||||
req.resp <- muxNewResponse{
|
||||
conn: proxyConn,
|
||||
err: err,
|
||||
}
|
||||
close(req.resp)
|
||||
case connID := <-m.channelClosed:
|
||||
if conn, ok := m.clients[string(connID[:])]; ok {
|
||||
conn.Detach(connID)
|
||||
delete(m.clients, string(connID[:]))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mux) Get(req *protocol.TelegramRequest) (*ProxyConn, error) {
|
||||
resp := make(chan muxNewResponse)
|
||||
m.channelNew <- muxNewRequest{
|
||||
req: req,
|
||||
resp: resp,
|
||||
}
|
||||
rv := <-resp
|
||||
|
||||
return rv.conn, rv.err
|
||||
}
|
||||
|
||||
func newMux(ctx context.Context) *mux {
|
||||
m := &mux{
|
||||
ctx: ctx,
|
||||
clients: make(map[string]*connection),
|
||||
channelClosed: make(chan conntypes.ConnID, 1),
|
||||
channelNew: make(chan muxNewRequest),
|
||||
}
|
||||
go m.run()
|
||||
|
||||
return m
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
"github.com/9seconds/mtg/mtproto/rpc"
|
||||
"github.com/9seconds/mtg/protocol"
|
||||
)
|
||||
|
||||
const (
|
||||
proxyConnWriteTimeout = 2 * time.Minute
|
||||
proxyConnReadTimeout = 2 * time.Minute
|
||||
)
|
||||
|
||||
type ProxyConn struct {
|
||||
closeOnce sync.Once
|
||||
req *protocol.TelegramRequest
|
||||
channelResponse chan *rpc.ProxyResponse
|
||||
channelClosed chan<- conntypes.ConnID
|
||||
channelWrite chan<- conntypes.Packet
|
||||
channelDone chan struct{}
|
||||
}
|
||||
|
||||
func (p *ProxyConn) Read() (*rpc.ProxyResponse, error) {
|
||||
timer := time.NewTimer(proxyConnReadTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-timer.C:
|
||||
return nil, ErrTimeout
|
||||
case <-p.channelDone:
|
||||
return nil, ErrClosed
|
||||
case packet := <-p.channelResponse:
|
||||
return packet, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ProxyConn) Write(packet conntypes.Packet) error {
|
||||
timer := time.NewTimer(proxyConnWriteTimeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case <-timer.C:
|
||||
return ErrTimeout
|
||||
case <-p.channelDone:
|
||||
return ErrClosed
|
||||
case p.channelWrite <- packet:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ProxyConn) put(response *rpc.ProxyResponse) {
|
||||
select {
|
||||
case <-p.channelDone:
|
||||
case p.channelResponse <- response:
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ProxyConn) Close() {
|
||||
p.closeOnce.Do(func() {
|
||||
close(p.channelDone)
|
||||
go func() {
|
||||
p.channelClosed <- p.req.ConnID
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
func newProxyConn(req *protocol.TelegramRequest, channelClosed chan<- conntypes.ConnID) *ProxyConn {
|
||||
return &ProxyConn{
|
||||
channelResponse: make(chan *rpc.ProxyResponse),
|
||||
channelDone: make(chan struct{}),
|
||||
channelClosed: channelClosed,
|
||||
req: req,
|
||||
}
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"github.com/9seconds/mtg/conntypes"
|
||||
)
|
||||
|
||||
type registry struct {
|
||||
conns map[string]*ctxChannel
|
||||
ctx context.Context
|
||||
mutex sync.RWMutex
|
||||
}
|
||||
|
||||
func (r *registry) Register(id conntypes.ConnID) ChannelReadCloser {
|
||||
channel := newCtxChannel(r.ctx)
|
||||
|
||||
r.mutex.Lock()
|
||||
r.conns[string(id[:])] = channel
|
||||
r.mutex.Unlock()
|
||||
|
||||
return channel
|
||||
}
|
||||
|
||||
func (r *registry) 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 *registry) getChannel(id conntypes.ConnID) (*ctxChannel, bool) {
|
||||
r.mutex.RLock()
|
||||
defer r.mutex.RUnlock()
|
||||
|
||||
if value, ok := r.conns[string(id[:])]; ok {
|
||||
return value, true
|
||||
}
|
||||
|
||||
return nil, false
|
||||
}
|
||||
Reference in New Issue
Block a user