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 14, 2023
1 parent 1578998 commit ee2bfb9
Showing 1 changed file with 66 additions and 0 deletions.
66 changes: 66 additions & 0 deletions pkg/adapters/grpc/traffic.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
package grpc

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

type DialerOption func(*dialerOptions)
type dialerOptions struct {
cm *route.ClusterManager
headers map[string]string
}

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 k, v := range options.headers {
tc.Headers[k] = v
}

instances, err := cm.Route(tc)
if err != nil {
return nil, err
}
if len(instances) == 0 {
return nil, errors.New("no matched provider")
}
return net.Dial("tcp", fmt.Sprintf("%s:%v", instances[0].Host, instances[0].Port))
}
}

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

func WithHeaders(headers map[string]string) DialerOption {
return func(d *dialerOptions) {
d.headers = headers
}
}

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

0 comments on commit ee2bfb9

Please sign in to comment.