-
-
Notifications
You must be signed in to change notification settings - Fork 126
/
pool.go
123 lines (104 loc) · 3 KB
/
pool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package telegram
import (
"context"
"go.uber.org/zap"
"golang.org/x/xerrors"
"github.com/gotd/td/internal/mtproto"
"github.com/gotd/td/internal/pool"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/telegram/internal/manager"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
)
// CloseInvoker is a closeable tg.Invoker.
type CloseInvoker interface {
tg.Invoker
Close() error
}
func (c *Client) createPool(dc int, max int64, creator func() pool.Conn) (*pool.DC, error) {
select {
case <-c.ctx.Done():
return nil, xerrors.Errorf("client already closed: %w", c.ctx.Err())
default:
}
p := pool.NewDC(c.ctx, dc, creator, pool.DCOptions{
Logger: c.log.Named("pool").With(zap.Int("dc_id", dc)),
MaxOpenConnections: max,
})
return p, nil
}
// Pool creates new multi-connection invoker to current DC.
func (c *Client) Pool(max int64) (CloseInvoker, error) {
if max < 0 {
return nil, xerrors.Errorf("invalid max value %d", max)
}
s := c.session.Load()
return c.createPool(s.DC, max, func() pool.Conn {
id := c.connsCounter.Inc()
return c.createConn(id, manager.ConnModeData, nil)
})
}
func (c *Client) dc(ctx context.Context, dcID int, max int64, dialer mtproto.Dialer) (*pool.DC, error) {
if max < 0 {
return nil, xerrors.Errorf("invalid max value %d", max)
}
dcList := dcs.FindDCs(c.cfg.Load().DCOptions, dcID, false)
if len(dcList) < 1 {
return nil, xerrors.Errorf("unknown DC %d", dcID)
}
c.log.Debug("Creating pool",
zap.Int("dc_id", dcID),
zap.Int64("max", max),
zap.Int("candidates", len(dcList)),
)
opts := c.opts
p, err := c.createPool(dcID, max, func() pool.Conn {
id := c.connsCounter.Inc()
c.sessionsMux.Lock()
session, ok := c.sessions[dcID]
if !ok {
session = pool.NewSyncSession(pool.Session{})
c.sessions[dcID] = session
}
c.sessionsMux.Unlock()
options, _ := session.Options(opts)
options.Logger = c.log.Named("conn").With(
zap.Int64("conn_id", id),
zap.Int("dc_id", dcID),
)
return c.create(
dialer, manager.ConnModeData, c.appID,
options, manager.ConnOptions{
DC: dcID,
Device: c.device,
Handler: c.asHandler(),
},
)
})
if err != nil {
return nil, xerrors.Errorf("create pool: %w", err)
}
_, err = c.transfer(ctx, tg.NewClient(p), dcID)
if err != nil {
// Ignore case then we are not authorized.
if unauthorized(err) {
return p, nil
}
// Kill pool if we got error.
_ = p.Close()
return nil, xerrors.Errorf("transfer: %w", err)
}
return p, nil
}
// DC creates new multi-connection invoker to given DC.
func (c *Client) DC(ctx context.Context, dc int, max int64) (CloseInvoker, error) {
return c.dc(ctx, dc, max, c.primaryDC(dc))
}
// MediaOnly creates new multi-connection invoker to given DC ID.
// It connects to MediaOnly DCs.
func (c *Client) MediaOnly(ctx context.Context, dc int, max int64) (CloseInvoker, error) {
return c.dc(ctx, dc, max, func(ctx context.Context) (transport.Conn, error) {
cfg := c.cfg.Load()
return c.resolver.MediaOnly(ctx, dc, cfg.DCOptions)
})
}