From b2da34baff3d802b9a3f955243baa32a641216cf Mon Sep 17 00:00:00 2001 From: Aleksandr Razumov Date: Mon, 29 Jan 2024 19:09:06 +0300 Subject: [PATCH 1/4] feat(clint): add OnTransfer callback Support external locking for auth transfer. --- telegram/client.go | 4 ++++ telegram/options.go | 11 +++++++++++ telegram/transfer.go | 33 +++++++++++++++++++++------------ 3 files changed, 36 insertions(+), 12 deletions(-) diff --git a/telegram/client.go b/telegram/client.go index 854ad993c3..643f222bea 100644 --- a/telegram/client.go +++ b/telegram/client.go @@ -128,6 +128,9 @@ type Client struct { // Tracing. tracer trace.Tracer + + // onTransfer is called on authorization transfer. + onTransfer func(ctx context.Context, tx func(ctx context.Context) error) error } // NewClient creates new unstarted client. @@ -161,6 +164,7 @@ func NewClient(appID int, appHash string, opt Options) *Client { migrationTimeout: opt.MigrationTimeout, noUpdatesMode: opt.NoUpdates, mw: opt.Middlewares, + onTransfer: opt.OnTransfer, } if opt.TracerProvider != nil { client.tracer = opt.TracerProvider.Tracer(oteltg.Name) diff --git a/telegram/options.go b/telegram/options.go index 0ac20a521d..3fedf177b8 100644 --- a/telegram/options.go +++ b/telegram/options.go @@ -94,6 +94,14 @@ type Options struct { // OpenTelemetry. TracerProvider trace.TracerProvider + + // OnTransfer is called on authorization transfer to acquire external lock. + // + // The function should call tx function to perform transfer, serializing it + // with external lock. + // + // The function must return error if tx function returned error. + OnTransfer func(ctx context.Context, tx func(context.Context) error) error } func (opt *Options) setDefaults() { @@ -140,6 +148,9 @@ func (opt *Options) setDefaults() { return nil }) } + if opt.OnTransfer == nil { + opt.OnTransfer = noopOnTransfer + } } func defaultBackoff(c clock.Clock) func() backoff.BackOff { diff --git a/telegram/transfer.go b/telegram/transfer.go index b6e7881d9f..f94a04ef83 100644 --- a/telegram/transfer.go +++ b/telegram/transfer.go @@ -17,20 +17,29 @@ func (c *Client) exportAuth(ctx context.Context, dcID int) (*tg.AuthExportedAuth return export, nil } +func noopOnTransfer(ctx context.Context, fn func(context.Context) error) error { return fn(ctx) } + // transfer exports current authorization and imports it to another DC. // See https://core.telegram.org/api/datacenter#authorization-transfer. func (c *Client) transfer(ctx context.Context, to *tg.Client, dc int) (tg.AuthAuthorizationClass, error) { - auth, err := c.exportAuth(ctx, dc) - if err != nil { - return nil, errors.Wrapf(err, "export to %d", dc) - } - - req := &tg.AuthImportAuthorizationRequest{} - req.FillFrom(auth) - r, err := to.AuthImportAuthorization(ctx, req) - if err != nil { - return nil, errors.Wrapf(err, "import from %d", dc) + var out tg.AuthAuthorizationClass + if err := c.onTransfer(ctx, func(ctx context.Context) error { + auth, err := c.exportAuth(ctx, dc) + if err != nil { + return errors.Wrapf(err, "export to %d", dc) + } + + req := &tg.AuthImportAuthorizationRequest{} + req.FillFrom(auth) + r, err := to.AuthImportAuthorization(ctx, req) + if err != nil { + return errors.Wrapf(err, "import from %d", dc) + } + + out = r + return nil + }); err != nil { + return nil, errors.Wrap(err, "onTransfer") } - - return r, nil + return out, nil } From 5333271bac3564cfc287259f25ee678cb36692a4 Mon Sep 17 00:00:00 2001 From: Aleksandr Razumov Date: Mon, 29 Jan 2024 19:11:27 +0300 Subject: [PATCH 2/4] feat(client): pass Client to OnTransfer --- telegram/client.go | 2 +- telegram/options.go | 2 +- telegram/transfer.go | 6 ++++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/telegram/client.go b/telegram/client.go index 643f222bea..564264b30a 100644 --- a/telegram/client.go +++ b/telegram/client.go @@ -130,7 +130,7 @@ type Client struct { tracer trace.Tracer // onTransfer is called on authorization transfer. - onTransfer func(ctx context.Context, tx func(ctx context.Context) error) error + onTransfer func(ctx context.Context, client *Client, tx func(ctx context.Context) error) error } // NewClient creates new unstarted client. diff --git a/telegram/options.go b/telegram/options.go index 3fedf177b8..724a581ed7 100644 --- a/telegram/options.go +++ b/telegram/options.go @@ -101,7 +101,7 @@ type Options struct { // with external lock. // // The function must return error if tx function returned error. - OnTransfer func(ctx context.Context, tx func(context.Context) error) error + OnTransfer func(ctx context.Context, c *Client, tx func(context.Context) error) error } func (opt *Options) setDefaults() { diff --git a/telegram/transfer.go b/telegram/transfer.go index f94a04ef83..f33a25cd63 100644 --- a/telegram/transfer.go +++ b/telegram/transfer.go @@ -17,13 +17,15 @@ func (c *Client) exportAuth(ctx context.Context, dcID int) (*tg.AuthExportedAuth return export, nil } -func noopOnTransfer(ctx context.Context, fn func(context.Context) error) error { return fn(ctx) } +func noopOnTransfer(ctx context.Context, _ *Client, fn func(context.Context) error) error { + return fn(ctx) +} // transfer exports current authorization and imports it to another DC. // See https://core.telegram.org/api/datacenter#authorization-transfer. func (c *Client) transfer(ctx context.Context, to *tg.Client, dc int) (tg.AuthAuthorizationClass, error) { var out tg.AuthAuthorizationClass - if err := c.onTransfer(ctx, func(ctx context.Context) error { + if err := c.onTransfer(ctx, c, func(ctx context.Context) error { auth, err := c.exportAuth(ctx, dc) if err != nil { return errors.Wrapf(err, "export to %d", dc) From ed3afc022d56b2f467c8fe0a4296c138db6b5d10 Mon Sep 17 00:00:00 2001 From: Aleksandr Razumov Date: Mon, 29 Jan 2024 19:12:46 +0300 Subject: [PATCH 3/4] test(client): set onTransfer --- telegram/client_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/telegram/client_test.go b/telegram/client_test.go index a2890d42ce..f6c8069cc9 100644 --- a/telegram/client_test.go +++ b/telegram/client_test.go @@ -86,6 +86,7 @@ func newTestClient(h testHandler) *Client { ctx: context.Background(), cancel: func() {}, updateHandler: UpdateHandlerFunc(func(ctx context.Context, u tg.UpdatesClass) error { return nil }), + onTransfer: noopOnTransfer, } client.init() From 615421bb99d10db2a0e0860f2bfe8063271e986a Mon Sep 17 00:00:00 2001 From: Aleksandr Razumov Date: Mon, 29 Jan 2024 19:20:28 +0300 Subject: [PATCH 4/4] refactor(client): add AuthTransferHandler type --- telegram/client.go | 4 ++-- telegram/options.go | 10 +++------- telegram/transfer.go | 8 ++++++++ 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/telegram/client.go b/telegram/client.go index 564264b30a..7476c001b3 100644 --- a/telegram/client.go +++ b/telegram/client.go @@ -129,8 +129,8 @@ type Client struct { // Tracing. tracer trace.Tracer - // onTransfer is called on authorization transfer. - onTransfer func(ctx context.Context, client *Client, tx func(ctx context.Context) error) error + // onTransfer is called in transfer. + onTransfer AuthTransferHandler } // NewClient creates new unstarted client. diff --git a/telegram/options.go b/telegram/options.go index 724a581ed7..af80cb06dd 100644 --- a/telegram/options.go +++ b/telegram/options.go @@ -95,13 +95,9 @@ type Options struct { // OpenTelemetry. TracerProvider trace.TracerProvider - // OnTransfer is called on authorization transfer to acquire external lock. - // - // The function should call tx function to perform transfer, serializing it - // with external lock. - // - // The function must return error if tx function returned error. - OnTransfer func(ctx context.Context, c *Client, tx func(context.Context) error) error + // OnTransfer is called during authorization transfer. + // See [AuthTransferHandler] for details. + OnTransfer AuthTransferHandler } func (opt *Options) setDefaults() { diff --git a/telegram/transfer.go b/telegram/transfer.go index f33a25cd63..ec77228348 100644 --- a/telegram/transfer.go +++ b/telegram/transfer.go @@ -17,6 +17,14 @@ func (c *Client) exportAuth(ctx context.Context, dcID int) (*tg.AuthExportedAuth return export, nil } +// AuthTransferHandler is a function that is called during authorization transfer. +// +// The fn callback should be serialized by user id via external locking. +// You can call [Client.Self] to acquire current user id. +// +// The fn callback must return fn error if any. +type AuthTransferHandler func(ctx context.Context, client *Client, fn func(context.Context) error) error + func noopOnTransfer(ctx context.Context, _ *Client, fn func(context.Context) error) error { return fn(ctx) }