Skip to content

Commit

Permalink
feat: adaptation for gRPC
Browse files Browse the repository at this point in the history
  • Loading branch information
saicaca committed Aug 23, 2023
1 parent 17f3e32 commit da12b40
Showing 1 changed file with 75 additions and 0 deletions.
75 changes: 75 additions & 0 deletions pkg/adapters/grpc/traffic.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
package grpc

import (
"context"
"errors"
"fmt"
"github.com/alibaba/sentinel-golang/core/route"
"github.com/alibaba/sentinel-golang/core/route/base"
"google.golang.org/grpc/metadata"
"net"
"strings"
)

type DialerOption func(*dialerOptions)
type dialerOptions struct {
cm *route.ClusterManager
md metadata.MD
}

func NewDialer(opts ...DialerOption) func(context.Context, string) (net.Conn, error) {
options := evaluateDialerOptions(opts)
return func(ctx context.Context, addr string) (net.Conn, error) {
cm := options.cm
parts := strings.Split(addr, "/")
if len(parts) != 2 {
return nil, errors.New("invalid address format")
}
tc := &base.TrafficContext{
ServiceName: parts[0],
MethodName: parts[1],
Headers: make(map[string]string),
}

for key, vals := range options.md {
if len(vals) == 0 {
continue
}
tc.Headers[key] = vals[len(vals)-1]
}

instance, err := cm.GetOne(tc)

for k, v := range tc.Baggage {
options.md[k] = []string{v}
}

if err != nil {
return nil, err
}
if instance == nil {
return nil, errors.New("no matched provider")
}
return net.Dial("tcp", fmt.Sprintf("%s:%v", instance.Host, instance.Port))
}
}

func WithClusterManager(cm *route.ClusterManager) DialerOption {
return func(d *dialerOptions) {
d.cm = cm
}
}

func WithMetadata(md metadata.MD) DialerOption {
return func(d *dialerOptions) {
d.md = md
}
}

func evaluateDialerOptions(opts []DialerOption) *dialerOptions {
ret := &dialerOptions{}
for _, opt := range opts {
opt(ret)
}
return ret
}

0 comments on commit da12b40

Please sign in to comment.