Merge pull request #393 from weibocom/dev

Dev

agent.go

@@ -13,6 +13,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -75,6 +76,7 @@ type Agent struct {
 	httpProxyServer   *mserver.HTTPProxyServer
 
 	manageHandlers map[string]http.Handler
+	envHandlers    map[string]map[string]http.Handler
 
 	svcLock      sync.Mutex
 	clsLock      sync.Mutex
@@ -109,6 +111,7 @@ func NewAgent(extfactory motan.ExtensionFactory) *Agent {
 	agent.agentPortServer = make(map[int]motan.Server)
 	agent.serviceRegistries = motan.NewCopyOnWriteMap()
 	agent.manageHandlers = make(map[string]http.Handler)
+	agent.envHandlers = make(map[string]map[string]http.Handler)
 	agent.serviceMap = motan.NewCopyOnWriteMap()
 	return agent
 }
@@ -778,9 +781,13 @@ func fillDefaultReqInfo(r motan.Request, url *motan.URL) {
 		}
 	} else {
 		if r.GetAttachment(mpro.MSource) == "" {
-			application := url.GetParam(motan.ApplicationKey, "")
-			if application != "" {
-				r.SetAttachment(mpro.MSource, application)
+			if app := r.GetAttachment(motan.ApplicationKey); app != "" {
+				r.SetAttachment(mpro.MSource, app)
+			} else {
+				application := url.GetParam(motan.ApplicationKey, "")
+				if application != "" {
+					r.SetAttachment(mpro.MSource, application)
+				}
 			}
 		}
 		if r.GetAttachment(mpro.MGroup) == "" {
@@ -1091,6 +1098,12 @@ func (a *Agent) RegisterManageHandler(path string, handler http.Handler) {
 	}
 }
 
+func (a *Agent) RegisterEnvHandlers(envStr string, handlers map[string]http.Handler) {
+	if envStr != "" && handlers != nil {
+		a.envHandlers[envStr] = handlers // override
+	}
+}
+
 func (a *Agent) startMServer() {
 	handlers := make(map[string]http.Handler, 16)
 	for k, v := range GetDefaultManageHandlers() {
@@ -1099,6 +1112,16 @@ func (a *Agent) startMServer() {
 	for k, v := range a.manageHandlers {
 		handlers[k] = v
 	}
+	// register env handlers
+	extHandelrs := os.Getenv(motan.HandlerEnvironmentName)
+	for _, k := range strings.Split(extHandelrs, ",") {
+		if v, ok := a.envHandlers[strings.TrimSpace(k)]; ok {
+			for kk, vv := range v {
+				handlers[kk] = vv
+			}
+
+		}
+	}
 	for k, v := range handlers {
 		a.mhandle(k, v)
 	}

agent_test.go

@@ -91,6 +91,58 @@ motan-refer:
 	assert.Equal(t, "Hello jack from motan server", resp.GetValue())
 	assert.Equal(t, 100, server.GetProcessPoolSize())
 }
+
+func Test_envHandler(t *testing.T) {
+	t.Parallel()
+	time.Sleep(time.Second * 3)
+	// start client mesh
+	ext := GetDefaultExtFactory()
+	os.Remove("agent.sock")
+	config, _ := config.NewConfigFromReader(bytes.NewReader([]byte(`
+motan-agent:
+  mport: 13500
+  port: 14821
+  eport: 14281
+  htport: 25282
+
+motan-registry:
+  direct:
+    protocol: direct
+    address: 127.0.0.1:22991
+
+motan-refer:
+    recom-engine-refer:
+      group: hello
+      path: helloService
+      protocol: motan2
+      registry: direct
+      asyncInitConnection: false
+      serialization: breeze`)))
+	agent := NewAgent(ext)
+	agent.RegisterEnvHandlers("testHandler", map[string]http.Handler{
+		"/test/test": testHandler(),
+	})
+	os.Setenv(core.HandlerEnvironmentName, "testHandler")
+	go agent.StartMotanAgentFromConfig(config)
+	time.Sleep(time.Second * 3)
+	client := http.Client{
+		Timeout: time.Second,
+	}
+	resp, err := client.Get("http://127.0.0.1:13500/test/test")
+	assert.Nil(t, err)
+	defer resp.Body.Close()
+	b, err := ioutil.ReadAll(resp.Body)
+	assert.Nil(t, err)
+	assert.Equal(t, string(b), "OK")
+	os.Unsetenv(core.HandlerEnvironmentName)
+}
+
+func testHandler() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("OK"))
+	}
+}
+
 func Test_unixClientCall2(t *testing.T) {
 	t.Parallel()
 	startServer(t, "helloService", 22992)

core/constants.go

@@ -70,6 +70,11 @@ const (
 	UnixSockProtocolFlag    = "unix://"
 )
 
+// metrics request application
+const (
+	MetricsReqApplication = "metricsReqApp"
+)
+
 // attachment keys
 const (
 	XForwardedForLower = "x-forwarded-for" // used as motan default proxy key
@@ -116,6 +121,9 @@ const (
 	GroupEnvironmentName     = "MESH_SERVICE_ADDITIONAL_GROUP"
 	DirectRPCEnvironmentName = "MESH_DIRECT_RPC"
 	FilterEnvironmentName    = "MESH_FILTERS"
+	HandlerEnvironmentName   = "MESH_ADMIN_EXT_HANDLERS"
+	RegGroupSuffix           = "RPC_REG_GROUP_SUFFIX"
+	SubGroupSuffix           = "MESH_MULTI_SUB_GROUP_SUFFIX"
 )
 
 // meta keys

core/globalContext.go

@@ -4,12 +4,11 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	cfg "github.com/weibocom/motan-go/config"
+	"github.com/weibocom/motan-go/log"
 	"os"
 	"reflect"
 	"strings"
-
-	cfg "github.com/weibocom/motan-go/config"
-	"github.com/weibocom/motan-go/log"
 )
 
 const (
@@ -69,7 +68,8 @@ var (
 	defaultConfigPath = "./"
 	defaultFileSuffix = ".yaml"
 
-	urlFields = map[string]bool{"protocol": true, "host": true, "port": true, "path": true, "group": true}
+	urlFields  = map[string]bool{"protocol": true, "host": true, "port": true, "path": true, "group": true}
+	extFilters = make(map[string]bool)
 )
 
 // all env flag in motan-go
@@ -87,6 +87,17 @@ var (
 	Recover     = flag.Bool("recover", false, "recover from accidental exit")
 )
 
+func AddRelevantFilter(filterStr string) {
+	k := strings.TrimSpace(filterStr)
+	if k != "" {
+		extFilters[k] = true
+	}
+}
+
+func GetRelevantFilters() map[string]bool {
+	return extFilters
+}
+
 func (c *Context) confToURLs(section string) map[string]*URL {
 	urls := map[string]*URL{}
 	sectionConf, _ := c.Config.GetSection(section)
@@ -238,6 +249,7 @@ func (c *Context) Initialize() {
 	c.parserBasicServices()
 	c.parseServices()
 	c.parseHTTPClients()
+	initSwitcher(c)
 }
 
 func (c *Context) parseSingleConfiguration() (*cfg.Config, error) {
@@ -397,11 +409,12 @@ func (c *Context) basicConfToURLs(section string) map[string]*URL {
 			newURL = url
 		}
 
-		//final filters: defaultFilter + globalFilter + filters + envFilter
+		//final filters: defaultFilter + globalFilter + filters + envFilter + relevantFilters
 		finalFilters := c.MergeFilterSet(
 			c.GetDefaultFilterSet(newURL),
 			c.GetGlobalFilterSet(newURL),
 			c.GetEnvGlobalFilterSet(),
+			GetRelevantFilters(),
 			c.GetFilterSet(newURL.GetStringParamsWithDefault(FilterKey, ""), ""),
 		)
 		if len(finalFilters) > 0 {
@@ -518,8 +531,57 @@ func (c *Context) parseMultipleServiceGroup(motanServiceMap map[string]*URL) {
 	}
 }
 
+func (c *Context) parseRegGroupSuffix(urlMap map[string]*URL) {
+	regGroupSuffix := os.Getenv(RegGroupSuffix)
+	if regGroupSuffix == "" {
+		return
+	}
+	filterMap := make(map[string]struct{}, len(urlMap))
+	for _, url := range urlMap {
+		filterMap[url.GetIdentityWithRegistry()] = struct{}{}
+	}
+	for k, url := range urlMap {
+		if strings.HasSuffix(url.Group, regGroupSuffix) {
+			continue
+		}
+		newUrl := url.Copy()
+		newUrl.Group += regGroupSuffix
+		if _, ok := filterMap[newUrl.GetIdentityWithRegistry()]; ok {
+			continue
+		}
+		filterMap[newUrl.GetIdentityWithRegistry()] = struct{}{}
+		urlMap[k] = newUrl
+	}
+}
+
+func (c *Context) parseSubGroupSuffix(urlMap map[string]*URL) {
+	subGroupSuffix := os.Getenv(SubGroupSuffix)
+	if subGroupSuffix == "" || c.AgentURL == nil {
+		return
+	}
+	filterMap := make(map[string]struct{}, len(urlMap)*2)
+	for _, url := range urlMap {
+		filterMap[url.GetIdentity()] = struct{}{}
+	}
+	for k, url := range urlMap {
+		if strings.HasSuffix(url.Group, subGroupSuffix) {
+			continue
+		}
+		groupWithSuffix := url.Group + subGroupSuffix
+		newUrl := url.Copy()
+		newUrl.Group += subGroupSuffix
+		if _, ok := filterMap[newUrl.GetIdentity()]; ok {
+			continue
+		}
+		filterMap[newUrl.GetIdentity()] = struct{}{}
+		urlMap["auto_"+k+groupWithSuffix] = newUrl
+	}
+}
+
 func (c *Context) parseRefers() {
-	c.RefersURLs = c.basicConfToURLs(refersSection)
+	referUrls := c.basicConfToURLs(refersSection)
+	c.parseSubGroupSuffix(referUrls)
+	c.RefersURLs = referUrls
 }
 
 func (c *Context) parseBasicRefers() {
@@ -529,6 +591,7 @@ func (c *Context) parseBasicRefers() {
 func (c *Context) parseServices() {
 	urlsMap := c.basicConfToURLs(servicesSection)
 	c.parseMultipleServiceGroup(urlsMap)
+	c.parseRegGroupSuffix(urlsMap)
 	c.ServiceURLs = urlsMap
 }