Use contexts for Conn wrapper

This commit is contained in:
9seconds
2018-07-28 13:21:57 +03:00
parent b0d86abc74
commit 9f20e8749a
9 changed files with 73 additions and 32 deletions
+5 -2
View File
@@ -1,6 +1,7 @@
package telegram
import (
"context"
"net"
"time"
@@ -38,12 +39,14 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
return conn, nil
}
func (t *tgDialer) dialRWC(addr, connID string) (wrappers.StreamReadWriteCloser, error) {
func (t *tgDialer) dialRWC(ctx context.Context, cancel context.CancelFunc,
addr, connID string) (wrappers.StreamReadWriteCloser, error) {
conn, err := t.dial(addr)
if err != nil {
return nil, err
}
tgConn := wrappers.NewConn(conn, connID, wrappers.ConnPurposeTelegram, t.conf.PublicIPv4, t.conf.PublicIPv6)
tgConn := wrappers.NewConn(ctx, cancel, conn, connID,
wrappers.ConnPurposeTelegram, t.conf.PublicIPv4, t.conf.PublicIPv6)
return tgConn, nil
}
+4 -2
View File
@@ -1,6 +1,7 @@
package telegram
import (
"context"
"net"
"github.com/juju/errors"
@@ -32,7 +33,8 @@ type directTelegram struct {
baseTelegram
}
func (t *directTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
func (t *directTelegram) Dial(ctx context.Context, cancel context.CancelFunc,
connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
dc := connOpts.DC
if dc < 0 {
dc = -dc
@@ -40,7 +42,7 @@ func (t *directTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (
dc = 1
}
return t.baseTelegram.dial(dc-1, connID, connOpts.ConnectionProto)
return t.baseTelegram.dial(ctx, cancel, dc-1, connID, connOpts.ConnectionProto)
}
func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts,
+3 -2
View File
@@ -2,6 +2,7 @@ package telegram
import (
"bufio"
"context"
"io/ioutil"
"net"
"net/http"
@@ -38,7 +39,7 @@ type middleTelegramCaller struct {
httpClient *http.Client
}
func (t *middleTelegramCaller) Dial(connID string,
func (t *middleTelegramCaller) Dial(ctx context.Context, cancel context.CancelFunc, connID string,
connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
dc := connOpts.DC
if dc == 0 {
@@ -47,7 +48,7 @@ func (t *middleTelegramCaller) Dial(connID string,
t.dialerMutex.RLock()
defer t.dialerMutex.RUnlock()
return t.baseTelegram.dial(dc, connID, connOpts.ConnectionProto)
return t.baseTelegram.dial(ctx, cancel, dc, connID, connOpts.ConnectionProto)
}
func (t *middleTelegramCaller) autoUpdate() {
+4 -3
View File
@@ -1,6 +1,7 @@
package telegram
import (
"context"
"math/rand"
"github.com/juju/errors"
@@ -11,7 +12,7 @@ import (
// Telegram is an interface for different Telegram work modes.
type Telegram interface {
Dial(string, *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error)
Dial(context.Context, context.CancelFunc, string, *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error)
Init(*mtproto.ConnectionOpts, wrappers.StreamReadWriteCloser) (wrappers.Wrap, error)
}
@@ -22,7 +23,7 @@ type baseTelegram struct {
v6Addresses map[int16][]string
}
func (b *baseTelegram) dial(dcIdx int16, connID string,
func (b *baseTelegram) dial(ctx context.Context, cancel context.CancelFunc, dcIdx int16, connID string,
proto mtproto.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
addrs := make([]string, 2)
@@ -38,7 +39,7 @@ func (b *baseTelegram) dial(dcIdx int16, connID string,
}
for _, addr := range addrs {
if conn, err := b.dialer.dialRWC(addr, connID); err == nil {
if conn, err := b.dialer.dialRWC(ctx, cancel, addr, connID); err == nil {
return conn, err
}
}