Merge pull request #363 from gotd/feat/client-middlewares

feat(telegram): reimplement middlewares

telegram/client.go

@@ -61,6 +61,10 @@ type clientConn interface {
 type Client struct {
 	// tg provides RPC calls via Client. Uses invoker below.
 	tg *tg.Client // immutable
+	// invoker implements tg.Invoker on top of Client and mw.
+	invoker tg.Invoker // immutable
+	// mw is list of middlewares used in invoker, can be blank.
+	mw []Middleware // immutable
 
 	// Telegram device information.
 	device DeviceConfig // immutable
@@ -227,7 +231,8 @@ func (c *Client) init() {
 	c.exported = make(chan *tg.AuthExportedAuthorization, 1)
 	c.sessions = map[int]*pool.SyncSession{}
 	c.subConns = map[int]CloseInvoker{}
-	c.tg = tg.NewClient(c)
+	c.invoker = chainMiddlewares(InvokeFunc(c.invokeDirect), c.mw...)
+	c.tg = tg.NewClient(c.invoker)
 }
 
 func (c *Client) restoreConnection(ctx context.Context) error {

telegram/invoke_raw.go → telegram/invoke.go

@@ -12,7 +12,7 @@ import (
 // Invoke invokes raw MTProto RPC method. It sends input and decodes result
 // into output.
 func (c *Client) Invoke(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
-	return c.invokeDirect(ctx, input, output)
+	return c.invoker.Invoke(ctx, input, output)
 }
 
 // invokeDirect directly invokes RPC method, automatically handling datacenter redirects.

telegram/middleware.go

@@ -0,0 +1,37 @@
+package telegram
+
+import (
+	"context"
+
+	"github.com/gotd/td/bin"
+	"github.com/gotd/td/tg"
+)
+
+// InvokeFunc implements tg.Invoker as function.
+type InvokeFunc func(ctx context.Context, input bin.Encoder, output bin.Decoder) error
+
+// Invoke implements tg.Invoker.
+func (i InvokeFunc) Invoke(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+	return i(ctx, input, output)
+}
+
+// Middleware returns new InvokeFunc for next invoker.
+type Middleware interface {
+	Handle(next tg.Invoker) InvokeFunc
+}
+
+// MiddlewareFunc implements Middleware as function.
+type MiddlewareFunc func(next tg.Invoker) InvokeFunc
+
+// Handle implements Middleware.
+func (m MiddlewareFunc) Handle(next tg.Invoker) InvokeFunc { return m(next) }
+
+func chainMiddlewares(invoker tg.Invoker, chain ...Middleware) tg.Invoker {
+	if len(chain) == 0 {
+		return invoker
+	}
+	for i := len(chain) - 1; i >= 0; i-- {
+		invoker = chain[i].Handle(invoker)
+	}
+	return invoker
+}

telegram/middleware_test.go

@@ -0,0 +1,37 @@
+package telegram
+
+import (
+	"context"
+	"testing"
+
+	"github.com/gotd/td/bin"
+	"github.com/gotd/td/tg"
+)
+
+type testPrintInvoker struct {
+	t *testing.T
+}
+
+func (t testPrintInvoker) Invoke(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+	t.t.Log("invoke")
+	return nil
+}
+
+func TestExampleMiddleware(t *testing.T) {
+	_ = chainMiddlewares(testPrintInvoker{t: t},
+		MiddlewareFunc(func(next tg.Invoker) InvokeFunc {
+			return func(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+				t.Log("First")
+				defer t.Log("After first")
+				return next.Invoke(ctx, input, output)
+			}
+		}),
+		MiddlewareFunc(func(next tg.Invoker) InvokeFunc {
+			return func(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
+				t.Log("Second")
+				defer t.Log("After second")
+				return next.Invoke(ctx, input, output)
+			}
+		}),
+	).Invoke(context.Background(), nil, nil)
+}

telegram/options.go

@@ -49,6 +49,9 @@ type Options struct {
 	SessionStorage SessionStorage
 	// UpdateHandler will be called on received update.
 	UpdateHandler UpdateHandler
+	// Middlewares list allows to wrap tg.Invoker. Can be useful for metrics,
+	// tracing, etc. Note that order is important.
+	Middlewares []Middleware
 
 	// AckBatchSize is maximum ack-s to buffer.
 	AckBatchSize int