fix(dcpool): add flood wait middleware

iyear · Dec 31, 2022 · fcb034b · fcb034b
1 parent 30e2bcf
commit fcb034b
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 94 deletions.
diff --git a/app/dl/dl.go b/app/dl/dl.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"github.com/fatih/color"
+	"github.com/gotd/contrib/middleware/floodwait"
 	"github.com/iyear/tdl/app/internal/tgc"
 	"github.com/iyear/tdl/pkg/consts"
 	"github.com/iyear/tdl/pkg/dcpool"
@@ -24,7 +25,7 @@ func Run(ctx context.Context, dir string, rewriteExt, skipSame bool, template st
 		color.Green("Preparing DC pool... It may take a while. size: %d", poolSize)
 
 		start := time.Now()
-		pool, err := dcpool.NewPool(ctx, c, poolSize)
+		pool, err := dcpool.NewPool(ctx, c, poolSize, floodwait.NewSimpleWaiter())
 		if err != nil {
 			return err
 		}

diff --git a/pkg/dcpool/dcpool.go b/pkg/dcpool/dcpool.go
@@ -14,18 +14,20 @@ var dcs = []int{1, 2, 3, 4, 5}
 
 type Pool interface {
 	Client(dc int) *tg.Client
-	Invoker(dc int) telegram.CloseInvoker
+	Invoker(dc int) tg.Invoker
 	Default() int
 	Close() error
 }
 
 type pool struct {
-	invokers map[int]telegram.CloseInvoker
+	invokers map[int]tg.Invoker
+	closes   map[int]func() error
 	_default int
 }
 
-func NewPool(ctx context.Context, c *telegram.Client, size int64) (Pool, error) {
-	m := make(map[int]telegram.CloseInvoker)
+func NewPool(ctx context.Context, c *telegram.Client, size int64, middlewares ...telegram.Middleware) (Pool, error) {
+	m := make(map[int]tg.Invoker)
+	closes := make(map[int]func() error)
 	mu := &sync.Mutex{}
 	curDC := c.Config().ThisDC
 
@@ -50,7 +52,8 @@ func NewPool(ctx context.Context, c *telegram.Client, size int64) (Pool, error)
 			}
 
 			mu.Lock()
-			m[dc] = invoker
+			closes[dc] = invoker.Close
+			m[dc] = chainMiddlewares(invoker, middlewares...)
 			mu.Unlock()
 
 			return nil
@@ -67,6 +70,7 @@ func NewPool(ctx context.Context, c *telegram.Client, size int64) (Pool, error)
 
 	return &pool{
 		invokers: m,
+		closes:   closes,
 		_default: curDC,
 	}, nil
 }
@@ -75,7 +79,7 @@ func (p *pool) Client(dc int) *tg.Client {
 	return tg.NewClient(p.Invoker(dc))
 }
 
-func (p *pool) Invoker(dc int) telegram.CloseInvoker {
+func (p *pool) Invoker(dc int) tg.Invoker {
 	i, ok := p.invokers[dc]
 	if !ok {
 		return p.invokers[p._default]
@@ -89,8 +93,8 @@ func (p *pool) Default() int {
 
 func (p *pool) Close() error {
 	var err error
-	for _, invokers := range p.invokers {
-		err = multierr.Append(err, invokers.Close())
+	for _, c := range p.closes {
+		err = multierr.Append(err, c())
 	}
 
 	return err

diff --git a/pkg/dcpool/lazypool.go b/pkg/dcpool/lazypool.go
diff --git a/pkg/dcpool/middlewares.go b/pkg/dcpool/middlewares.go
@@ -0,0 +1,17 @@
+package dcpool
+
+import (
+	"github.com/gotd/td/telegram"
+	"github.com/gotd/td/tg"
+)
+
+func chainMiddlewares(invoker tg.Invoker, chain ...telegram.Middleware) tg.Invoker {
+	if len(chain) == 0 {
+		return invoker
+	}
+	for i := len(chain) - 1; i >= 0; i-- {
+		invoker = chain[i].Handle(invoker)
+	}
+
+	return invoker
+}