Add CloudWeGo Kitex adapter (#473)

Author: li-jin-gou

pkg/adapters/kitex/client.go

@@ -0,0 +1,36 @@
+package kitex
+
+import (
+	"context"
+
+	sentinel "github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/base"
+	"github.com/cloudwego/kitex/pkg/endpoint"
+)
+
+// SentinelClientMiddleware returns new client.Middleware
+// Default resource name is {service's name}:{method}
+// Default block fallback is returning blockError
+// Define your own behavior by setting serverOptions
+func SentinelClientMiddleware(opts ...Option) func(endpoint.Endpoint) endpoint.Endpoint {
+	options := newOptions(opts)
+	return func(next endpoint.Endpoint) endpoint.Endpoint {
+		return func(ctx context.Context, req, resp interface{}) error {
+			resourceName := options.ResourceExtract(ctx, req, resp)
+			entry, blockErr := sentinel.Entry(
+				resourceName,
+				sentinel.WithResourceType(base.ResTypeRPC),
+				sentinel.WithTrafficType(base.Outbound),
+			)
+			if blockErr != nil {
+				return options.BlockFallback(ctx, req, resp, blockErr)
+			}
+			defer entry.Exit()
+			err := next(ctx, req, resp)
+			if err != nil {
+				sentinel.TraceError(entry, err)
+			}
+			return err
+		}
+	}
+}

pkg/adapters/kitex/client_test.go

@@ -0,0 +1,52 @@
+package kitex
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	sentinel "github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/flow"
+	"github.com/cloudwego/kitex-examples/hello/kitex_gen/api"
+	"github.com/cloudwego/kitex-examples/hello/kitex_gen/api/hello"
+	"github.com/cloudwego/kitex/client"
+	"github.com/stretchr/testify/assert"
+)
+
+const FakeErrorMsg = "fake error for testing"
+
+func TestSentinelClientMiddleware(t *testing.T) {
+	bf := func(ctx context.Context, req, resp interface{}, blockErr error) error {
+		return errors.New(FakeErrorMsg)
+	}
+	c, err := hello.NewClient("hello",
+		client.WithMiddleware(SentinelClientMiddleware(
+			WithBlockFallback(bf))))
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = sentinel.InitDefault()
+	if err != nil {
+		t.Fatal(err)
+	}
+	req := &api.Request{}
+	t.Run("success", func(t *testing.T) {
+		_, err := flow.LoadRules([]*flow.Rule{
+			{
+				Resource:               "hello:echo",
+				Threshold:              1.0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
+			},
+		})
+		assert.Nil(t, err)
+		_, err = c.Echo(context.Background(), req)
+		assert.NotNil(t, err)
+		assert.NotEqual(t, FakeErrorMsg, err.Error())
+		t.Run("second fail", func(t *testing.T) {
+			_, err = c.Echo(context.Background(), req)
+			assert.NotNil(t, err)
+			assert.Equal(t, FakeErrorMsg, err.Error())
+		})
+	})
+}

pkg/adapters/kitex/doc.go

@@ -0,0 +1,17 @@
+/**
+This package provides Sentinel integration for Kitex.
+
+For server side, users may append a Sentinel middleware to Kitex service, like:
+		import (
+			sentinelPlugin "github.com/alibaba/sentinel-golang/pkg/adapters/kitex"
+		)
+		srv := hello.NewServer(new(HelloImpl),server.WithMiddleware(SentinelServerMiddleware()))
+The plugin extracts service name and service method as the resource name by default.
+Users may provide customized resource name extractor when creating new
+Sentinel middleware (via WithResourceExtract options).
+
+Fallback logic: the plugin will return the BlockError by default
+if current request is blocked by Sentinel rules. Users may also
+provide customized fallback logic via WithBlockFallback(handler) options.
+*/
+package kitex

pkg/adapters/kitex/go.mod

@@ -0,0 +1,10 @@
+module github.com/alibaba/sentinel-golang/pkg/adapters/kitex
+
+go 1.16
+
+require (
+	github.com/alibaba/sentinel-golang v1.0.4
+	github.com/cloudwego/kitex v0.3.4
+	github.com/cloudwego/kitex-examples v0.1.0
+	github.com/stretchr/testify v1.7.0
+)

pkg/adapters/kitex/go.sum

@@ -0,0 +1,54 @@
+package kitex
+
+import (
+	"context"
+
+	"github.com/cloudwego/kitex/pkg/rpcinfo"
+)
+
+type Option struct {
+	F func(o *options)
+}
+
+type options struct {
+	ResourceExtract func(ctx context.Context, req, resp interface{}) string
+	BlockFallback   func(ctx context.Context, req, resp interface{}, blockErr error) error
+}
+
+func DefaultBlockFallback(ctx context.Context, req, resp interface{}, blockErr error) error {
+	return blockErr
+}
+
+func DefaultResourceExtract(ctx context.Context, req, resp interface{}) string {
+	ri := rpcinfo.GetRPCInfo(ctx)
+	return ri.To().ServiceName() + ":" + ri.To().Method()
+}
+
+func newOptions(opts []Option) *options {
+	o := &options{
+		ResourceExtract: DefaultResourceExtract,
+		BlockFallback:   DefaultBlockFallback,
+	}
+	o.Apply(opts)
+	return o
+}
+
+func (o *options) Apply(opts []Option) {
+	for _, op := range opts {
+		op.F(o)
+	}
+}
+
+// WithResourceExtract sets the resource extractor
+func WithResourceExtract(f func(ctx context.Context, req, resp interface{}) string) Option {
+	return Option{F: func(o *options) {
+		o.ResourceExtract = f
+	}}
+}
+
+// WithBlockFallback sets the fallback handler
+func WithBlockFallback(f func(ctx context.Context, req, resp interface{}, blockErr error) error) Option {
+	return Option{func(o *options) {
+		o.BlockFallback = f
+	}}
+}

pkg/adapters/kitex/options.go

@@ -0,0 +1,36 @@
+package kitex
+
+import (
+	"context"
+
+	sentinel "github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/base"
+	"github.com/cloudwego/kitex/pkg/endpoint"
+)
+
+// SentinelServerMiddleware returns new server.Middleware
+// Default resource name is {service's name}:{method}
+// Default block fallback is returning blockError
+// Define your own behavior by setting serverOptions
+func SentinelServerMiddleware(opts ...Option) func(endpoint.Endpoint) endpoint.Endpoint {
+	options := newOptions(opts)
+	return func(next endpoint.Endpoint) endpoint.Endpoint {
+		return func(ctx context.Context, req, resp interface{}) error {
+			resourceName := options.ResourceExtract(ctx, req, resp)
+			entry, blockErr := sentinel.Entry(
+				resourceName,
+				sentinel.WithResourceType(base.ResTypeRPC),
+				sentinel.WithTrafficType(base.Inbound),
+			)
+			if blockErr != nil {
+				return options.BlockFallback(ctx, req, resp, blockErr)
+			}
+			defer entry.Exit()
+			err := next(ctx, req, resp)
+			if err != nil {
+				sentinel.TraceError(entry, err)
+			}
+			return err
+		}
+	}
+}

pkg/adapters/kitex/server.go

@@ -0,0 +1,67 @@
+package kitex
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"testing"
+	"time"
+
+	sentinel "github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/flow"
+	"github.com/cloudwego/kitex-examples/hello/kitex_gen/api"
+	"github.com/cloudwego/kitex-examples/hello/kitex_gen/api/hello"
+	"github.com/cloudwego/kitex/client"
+	"github.com/cloudwego/kitex/pkg/rpcinfo"
+	"github.com/cloudwego/kitex/server"
+	"github.com/stretchr/testify/assert"
+)
+
+// HelloImpl implements the last service interface defined in the IDL.
+type HelloImpl struct{}
+
+// Echo implements the HelloImpl interface.
+func (s *HelloImpl) Echo(ctx context.Context, req *api.Request) (resp *api.Response, err error) {
+	resp = &api.Response{Message: req.Message}
+	return
+}
+
+func TestSentinelServerMiddleware(t *testing.T) {
+	bf := func(ctx context.Context, req, resp interface{}, blockErr error) error {
+		return errors.New(FakeErrorMsg)
+	}
+	srv := hello.NewServer(new(HelloImpl),
+		server.WithServerBasicInfo(&rpcinfo.EndpointBasicInfo{ServiceName: "hello"}),
+		server.WithMiddleware(SentinelServerMiddleware(
+			WithBlockFallback(bf),
+		)))
+	go srv.Run()
+	defer srv.Stop()
+	time.Sleep(1 * time.Second)
+
+	c, err := hello.NewClient("hello", client.WithHostPorts(":8888"))
+	assert.Nil(t, err)
+
+	err = sentinel.InitDefault()
+	assert.Nil(t, err)
+	req := &api.Request{}
+	t.Run("success", func(t *testing.T) {
+		_, err = flow.LoadRules([]*flow.Rule{
+			{
+				Resource:               "hello:echo",
+				Threshold:              1.0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
+			},
+		})
+		assert.Nil(t, err)
+		_, err := c.Echo(context.TODO(), req)
+		assert.Nil(t, err)
+
+		t.Run("second fail", func(t *testing.T) {
+			_, err := c.Echo(context.TODO(), req)
+			assert.Error(t, err)
+			assert.True(t, strings.Contains(err.Error(), FakeErrorMsg))
+		})
+	})
+}