Skip to content

Commit a6f5d1c

Browse files
rayzhang0603rayzhang0603
authored
Merge pull request #393 from weibocom/dev
Dev
2 parents db996c8 + 4f95c41 commit a6f5d1c

20 files changed

+716
-45
lines changed

‎agent.go

+26-3
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"path/filepath"
1414
"runtime"
1515
"strconv"
16+
"strings"
1617
"sync"
1718
"sync/atomic"
1819
"time"
@@ -75,6 +76,7 @@ type Agent struct {
7576
httpProxyServer *mserver.HTTPProxyServer
7677

7778
manageHandlers map[string]http.Handler
79+
envHandlers map[string]map[string]http.Handler
7880

7981
svcLock sync.Mutex
8082
clsLock sync.Mutex
@@ -109,6 +111,7 @@ func NewAgent(extfactory motan.ExtensionFactory) *Agent {
109111
agent.agentPortServer = make(map[int]motan.Server)
110112
agent.serviceRegistries = motan.NewCopyOnWriteMap()
111113
agent.manageHandlers = make(map[string]http.Handler)
114+
agent.envHandlers = make(map[string]map[string]http.Handler)
112115
agent.serviceMap = motan.NewCopyOnWriteMap()
113116
return agent
114117
}
@@ -778,9 +781,13 @@ func fillDefaultReqInfo(r motan.Request, url *motan.URL) {
778781
}
779782
} else {
780783
if r.GetAttachment(mpro.MSource) == "" {
781-
application := url.GetParam(motan.ApplicationKey, "")
782-
if application != "" {
783-
r.SetAttachment(mpro.MSource, application)
784+
if app := r.GetAttachment(motan.ApplicationKey); app != "" {
785+
r.SetAttachment(mpro.MSource, app)
786+
} else {
787+
application := url.GetParam(motan.ApplicationKey, "")
788+
if application != "" {
789+
r.SetAttachment(mpro.MSource, application)
790+
}
784791
}
785792
}
786793
if r.GetAttachment(mpro.MGroup) == "" {
@@ -1091,6 +1098,12 @@ func (a *Agent) RegisterManageHandler(path string, handler http.Handler) {
10911098
}
10921099
}
10931100

1101+
func (a *Agent) RegisterEnvHandlers(envStr string, handlers map[string]http.Handler) {
1102+
if envStr != "" && handlers != nil {
1103+
a.envHandlers[envStr] = handlers // override
1104+
}
1105+
}
1106+
10941107
func (a *Agent) startMServer() {
10951108
handlers := make(map[string]http.Handler, 16)
10961109
for k, v := range GetDefaultManageHandlers() {
@@ -1099,6 +1112,16 @@ func (a *Agent) startMServer() {
10991112
for k, v := range a.manageHandlers {
11001113
handlers[k] = v
11011114
}
1115+
// register env handlers
1116+
extHandelrs := os.Getenv(motan.HandlerEnvironmentName)
1117+
for _, k := range strings.Split(extHandelrs, ",") {
1118+
if v, ok := a.envHandlers[strings.TrimSpace(k)]; ok {
1119+
for kk, vv := range v {
1120+
handlers[kk] = vv
1121+
}
1122+
1123+
}
1124+
}
11021125
for k, v := range handlers {
11031126
a.mhandle(k, v)
11041127
}

‎agent_test.go

+52
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,58 @@ motan-refer:
9191
assert.Equal(t, "Hello jack from motan server", resp.GetValue())
9292
assert.Equal(t, 100, server.GetProcessPoolSize())
9393
}
94+
95+
func Test_envHandler(t *testing.T) {
96+
t.Parallel()
97+
time.Sleep(time.Second * 3)
98+
// start client mesh
99+
ext := GetDefaultExtFactory()
100+
os.Remove("agent.sock")
101+
config, _ := config.NewConfigFromReader(bytes.NewReader([]byte(`
102+
motan-agent:
103+
mport: 13500
104+
port: 14821
105+
eport: 14281
106+
htport: 25282
107+
108+
motan-registry:
109+
direct:
110+
protocol: direct
111+
address: 127.0.0.1:22991
112+
113+
motan-refer:
114+
recom-engine-refer:
115+
group: hello
116+
path: helloService
117+
protocol: motan2
118+
registry: direct
119+
asyncInitConnection: false
120+
serialization: breeze`)))
121+
agent := NewAgent(ext)
122+
agent.RegisterEnvHandlers("testHandler", map[string]http.Handler{
123+
"/test/test": testHandler(),
124+
})
125+
os.Setenv(core.HandlerEnvironmentName, "testHandler")
126+
go agent.StartMotanAgentFromConfig(config)
127+
time.Sleep(time.Second * 3)
128+
client := http.Client{
129+
Timeout: time.Second,
130+
}
131+
resp, err := client.Get("http://127.0.0.1:13500/test/test")
132+
assert.Nil(t, err)
133+
defer resp.Body.Close()
134+
b, err := ioutil.ReadAll(resp.Body)
135+
assert.Nil(t, err)
136+
assert.Equal(t, string(b), "OK")
137+
os.Unsetenv(core.HandlerEnvironmentName)
138+
}
139+
140+
func testHandler() http.HandlerFunc {
141+
return func(w http.ResponseWriter, r *http.Request) {
142+
w.Write([]byte("OK"))
143+
}
144+
}
145+
94146
func Test_unixClientCall2(t *testing.T) {
95147
t.Parallel()
96148
startServer(t, "helloService", 22992)

‎core/constants.go

+8
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,11 @@ const (
7070
UnixSockProtocolFlag = "unix://"
7171
)
7272

73+
// metrics request application
74+
const (
75+
MetricsReqApplication = "metricsReqApp"
76+
)
77+
7378
// attachment keys
7479
const (
7580
XForwardedForLower = "x-forwarded-for" // used as motan default proxy key
@@ -116,6 +121,9 @@ const (
116121
GroupEnvironmentName = "MESH_SERVICE_ADDITIONAL_GROUP"
117122
DirectRPCEnvironmentName = "MESH_DIRECT_RPC"
118123
FilterEnvironmentName = "MESH_FILTERS"
124+
HandlerEnvironmentName = "MESH_ADMIN_EXT_HANDLERS"
125+
RegGroupSuffix = "RPC_REG_GROUP_SUFFIX"
126+
SubGroupSuffix = "MESH_MULTI_SUB_GROUP_SUFFIX"
119127
)
120128

121129
// meta keys

‎core/globalContext.go

+69-6
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,11 @@ import (
44
"errors"
55
"flag"
66
"fmt"
7+
cfg "github.com/weibocom/motan-go/config"
8+
"github.com/weibocom/motan-go/log"
79
"os"
810
"reflect"
911
"strings"
10-
11-
cfg "github.com/weibocom/motan-go/config"
12-
"github.com/weibocom/motan-go/log"
1312
)
1413

1514
const (
@@ -69,7 +68,8 @@ var (
6968
defaultConfigPath = "./"
7069
defaultFileSuffix = ".yaml"
7170

72-
urlFields = map[string]bool{"protocol": true, "host": true, "port": true, "path": true, "group": true}
71+
urlFields = map[string]bool{"protocol": true, "host": true, "port": true, "path": true, "group": true}
72+
extFilters = make(map[string]bool)
7373
)
7474

7575
// all env flag in motan-go
@@ -87,6 +87,17 @@ var (
8787
Recover = flag.Bool("recover", false, "recover from accidental exit")
8888
)
8989

90+
func AddRelevantFilter(filterStr string) {
91+
k := strings.TrimSpace(filterStr)
92+
if k != "" {
93+
extFilters[k] = true
94+
}
95+
}
96+
97+
func GetRelevantFilters() map[string]bool {
98+
return extFilters
99+
}
100+
90101
func (c *Context) confToURLs(section string) map[string]*URL {
91102
urls := map[string]*URL{}
92103
sectionConf, _ := c.Config.GetSection(section)
@@ -238,6 +249,7 @@ func (c *Context) Initialize() {
238249
c.parserBasicServices()
239250
c.parseServices()
240251
c.parseHTTPClients()
252+
initSwitcher(c)
241253
}
242254

243255
func (c *Context) parseSingleConfiguration() (*cfg.Config, error) {
@@ -397,11 +409,12 @@ func (c *Context) basicConfToURLs(section string) map[string]*URL {
397409
newURL = url
398410
}
399411

400-
//final filters: defaultFilter + globalFilter + filters + envFilter
412+
//final filters: defaultFilter + globalFilter + filters + envFilter + relevantFilters
401413
finalFilters := c.MergeFilterSet(
402414
c.GetDefaultFilterSet(newURL),
403415
c.GetGlobalFilterSet(newURL),
404416
c.GetEnvGlobalFilterSet(),
417+
GetRelevantFilters(),
405418
c.GetFilterSet(newURL.GetStringParamsWithDefault(FilterKey, ""), ""),
406419
)
407420
if len(finalFilters) > 0 {
@@ -518,8 +531,57 @@ func (c *Context) parseMultipleServiceGroup(motanServiceMap map[string]*URL) {
518531
}
519532
}
520533

534+
func (c *Context) parseRegGroupSuffix(urlMap map[string]*URL) {
535+
regGroupSuffix := os.Getenv(RegGroupSuffix)
536+
if regGroupSuffix == "" {
537+
return
538+
}
539+
filterMap := make(map[string]struct{}, len(urlMap))
540+
for _, url := range urlMap {
541+
filterMap[url.GetIdentityWithRegistry()] = struct{}{}
542+
}
543+
for k, url := range urlMap {
544+
if strings.HasSuffix(url.Group, regGroupSuffix) {
545+
continue
546+
}
547+
newUrl := url.Copy()
548+
newUrl.Group += regGroupSuffix
549+
if _, ok := filterMap[newUrl.GetIdentityWithRegistry()]; ok {
550+
continue
551+
}
552+
filterMap[newUrl.GetIdentityWithRegistry()] = struct{}{}
553+
urlMap[k] = newUrl
554+
}
555+
}
556+
557+
func (c *Context) parseSubGroupSuffix(urlMap map[string]*URL) {
558+
subGroupSuffix := os.Getenv(SubGroupSuffix)
559+
if subGroupSuffix == "" || c.AgentURL == nil {
560+
return
561+
}
562+
filterMap := make(map[string]struct{}, len(urlMap)*2)
563+
for _, url := range urlMap {
564+
filterMap[url.GetIdentity()] = struct{}{}
565+
}
566+
for k, url := range urlMap {
567+
if strings.HasSuffix(url.Group, subGroupSuffix) {
568+
continue
569+
}
570+
groupWithSuffix := url.Group + subGroupSuffix
571+
newUrl := url.Copy()
572+
newUrl.Group += subGroupSuffix
573+
if _, ok := filterMap[newUrl.GetIdentity()]; ok {
574+
continue
575+
}
576+
filterMap[newUrl.GetIdentity()] = struct{}{}
577+
urlMap["auto_"+k+groupWithSuffix] = newUrl
578+
}
579+
}
580+
521581
func (c *Context) parseRefers() {
522-
c.RefersURLs = c.basicConfToURLs(refersSection)
582+
referUrls := c.basicConfToURLs(refersSection)
583+
c.parseSubGroupSuffix(referUrls)
584+
c.RefersURLs = referUrls
523585
}
524586

525587
func (c *Context) parseBasicRefers() {
@@ -529,6 +591,7 @@ func (c *Context) parseBasicRefers() {
529591
func (c *Context) parseServices() {
530592
urlsMap := c.basicConfToURLs(servicesSection)
531593
c.parseMultipleServiceGroup(urlsMap)
594+
c.parseRegGroupSuffix(urlsMap)
532595
c.ServiceURLs = urlsMap
533596
}
534597

0 commit comments

Comments
 (0)