From b6bedcd5225b5522659b7e07ff79513617fe85c1 Mon Sep 17 00:00:00 2001
From: xiandong <xiandong@italki.com>
Date: Fri, 17 Feb 2023 12:58:31 +0800
Subject: [PATCH] add kindJaegerUdp

---
 core/trace/agent.go      | 18 +++++++++---------
 core/trace/agent_test.go | 34 +++++++++-------------------------
 core/trace/config.go     | 31 ++++++++++++++++++-------------
 3 files changed, 36 insertions(+), 47 deletions(-)

diff --git a/core/trace/agent.go b/core/trace/agent.go
index 0958ea5f..d8749912 100644
--- a/core/trace/agent.go
+++ b/core/trace/agent.go
@@ -18,10 +18,11 @@ import (
 )
 
 const (
-	kindJaeger   = "jaeger"
-	kindZipkin   = "zipkin"
-	kindOtlpGrpc = "otlpgrpc"
-	kindOtlpHttp = "otlphttp"
+	kindJaeger    = "jaeger"
+	kindJaegerUdp = "jaegerudp"
+	kindZipkin    = "zipkin"
+	kindOtlpGrpc  = "otlpgrpc"
+	kindOtlpHttp  = "otlphttp"
 )
 
 var (
@@ -35,7 +36,7 @@ func StartAgent(c Config) {
 	lock.Lock()
 	defer lock.Unlock()
 
-	_, ok := agents[c.getEndpoint()]
+	_, ok := agents[c.Endpoint]
 	if ok {
 		return
 	}
@@ -45,7 +46,7 @@ func StartAgent(c Config) {
 		return
 	}
 
-	agents[c.getEndpoint()] = lang.Placeholder
+	agents[c.Endpoint] = lang.Placeholder
 }
 
 // StopAgent shuts down the span processors in the order they were registered.
@@ -57,10 +58,9 @@ func createExporter(c Config) (sdktrace.SpanExporter, error) {
 	// Just support jaeger and zipkin now, more for later
 	switch c.Batcher {
 	case kindJaeger:
-		if c.isAgentEndPoint() {
-			return jaeger.New(jaeger.WithAgentEndpoint(jaeger.WithAgentHost(c.AgentHost), jaeger.WithAgentPort(c.AgentPort)))
-		}
 		return jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.Endpoint)))
+	case kindJaegerUdp:
+		return jaeger.New(jaeger.WithAgentEndpoint(jaeger.WithAgentHost(c.getEndpointHost()), jaeger.WithAgentPort(c.getEndpointPort())))
 	case kindZipkin:
 		return zipkin.New(c.Endpoint)
 	case kindOtlpGrpc:
diff --git a/core/trace/agent_test.go b/core/trace/agent_test.go
index dc9172b4..e77bd578 100644
--- a/core/trace/agent_test.go
+++ b/core/trace/agent_test.go
@@ -11,12 +11,11 @@ func TestStartAgent(t *testing.T) {
 	logx.Disable()
 
 	const (
-		endpoint1  = "localhost:1234"
-		endpoint2  = "remotehost:1234"
-		endpoint3  = "localhost:1235"
-		endpoint4  = "localhost:1236"
-		agentHost1 = "localhost"
-		agentPort1 = "6831"
+		endpoint1 = "localhost:1234"
+		endpoint2 = "remotehost:1234"
+		endpoint3 = "localhost:1235"
+		endpoint4 = "localhost:1236"
+		endpoint5 = "localhost:6831"
 	)
 	c1 := Config{
 		Name: "foo",
@@ -47,17 +46,9 @@ func TestStartAgent(t *testing.T) {
 		Batcher:  kindOtlpHttp,
 	}
 	c7 := Config{
-		Name:      "jaegerUDP",
-		AgentHost: agentHost1,
-		AgentPort: agentPort1,
-		Batcher:   kindJaeger,
-	}
-	c8 := Config{
-		Name:      "jaegerUDP",
-		AgentHost: agentHost1,
-		AgentPort: agentPort1,
-		Endpoint:  endpoint1,
-		Batcher:   kindJaeger,
+		Name:     "UDP",
+		Endpoint: endpoint5,
+		Batcher:  kindJaegerUdp,
 	}
 
 	StartAgent(c1)
@@ -68,7 +59,6 @@ func TestStartAgent(t *testing.T) {
 	StartAgent(c5)
 	StartAgent(c6)
 	StartAgent(c7)
-	StartAgent(c8)
 
 	lock.Lock()
 	defer lock.Unlock()
@@ -81,12 +71,6 @@ func TestStartAgent(t *testing.T) {
 	assert.True(t, ok)
 	_, ok = agents[endpoint2]
 	assert.False(t, ok)
-	_, ok = agents[c2.getEndpoint()]
-	assert.True(t, ok)
-	_, ok = agents[c3.getEndpoint()]
-	assert.False(t, ok)
-	_, ok = agents[c7.getEndpoint()]
-	assert.True(t, ok)
-	_, ok = agents[c8.getEndpoint()]
+	_, ok = agents[endpoint5]
 	assert.True(t, ok)
 }
diff --git a/core/trace/config.go b/core/trace/config.go
index 4abf2e8b..5b181202 100644
--- a/core/trace/config.go
+++ b/core/trace/config.go
@@ -1,27 +1,32 @@
 package trace
 
-import "fmt"
+import (
+	"strings"
+)
 
 // TraceName represents the tracing name.
 const TraceName = "go-zero"
 
 // A Config is an opentelemetry config.
 type Config struct {
-	Name      string  `json:",optional"`
-	AgentHost string  `json:",optional"`
-	AgentPort string  `json:",optional"`
-	Endpoint  string  `json:",optional"`
-	Sampler   float64 `json:",default=1.0"`
-	Batcher   string  `json:",default=jaeger,options=jaeger|jaegerudp|zipkin|otlpgrpc|otlphttp"`
+	Name     string  `json:",optional"`
+	Endpoint string  `json:",optional"`
+	Sampler  float64 `json:",default=1.0"`
+	Batcher  string  `json:",default=jaeger,options=jaeger|jaegerudp|zipkin|otlpgrpc|otlphttp"`
 }
 
-func (c *Config) isAgentEndPoint() bool {
-	return len(c.AgentHost) != 0 && len(c.AgentPort) != 0
+func (c *Config) getEndpointHost() string {
+	EndpointSlice := strings.Split(c.Endpoint, ":")
+	if len(EndpointSlice) > 0 {
+		return strings.TrimSpace(EndpointSlice[0])
+	}
+	return ""
 }
 
-func (c *Config) getEndpoint() string {
-	if c.isAgentEndPoint() {
-		return fmt.Sprintf("%s:%s", c.AgentHost, c.AgentPort)
+func (c *Config) getEndpointPort() string {
+	EndpointSlice := strings.Split(c.Endpoint, ":")
+	if len(EndpointSlice) > 1 {
+		return strings.TrimSpace(EndpointSlice[1])
 	}
-	return c.Endpoint
+	return ""
 }