From 12060c9c0c4a5a7c1fa093d53bb79c1afafc12a3 Mon Sep 17 00:00:00 2001
From: xiandong <xiandong@italki.com>
Date: Fri, 17 Feb 2023 12:41:26 +0800
Subject: [PATCH] opentelemetry support AgentHost, AgentPort

---
 core/trace/agent.go      |  9 ++++-----
 core/trace/agent_test.go | 35 ++++++++++++++++++++++++++++++-----
 core/trace/config.go     | 15 ++++++++++++++-
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/core/trace/agent.go b/core/trace/agent.go
index 053f9e3e..0958ea5f 100644
--- a/core/trace/agent.go
+++ b/core/trace/agent.go
@@ -35,7 +35,7 @@ func StartAgent(c Config) {
 	lock.Lock()
 	defer lock.Unlock()
 
-	_, ok := agents[c.Endpoint]
+	_, ok := agents[c.getEndpoint()]
 	if ok {
 		return
 	}
@@ -45,7 +45,7 @@ func StartAgent(c Config) {
 		return
 	}
 
-	agents[c.Endpoint] = lang.Placeholder
+	agents[c.getEndpoint()] = lang.Placeholder
 }
 
 // StopAgent shuts down the span processors in the order they were registered.
@@ -57,11 +57,10 @@ func createExporter(c Config) (sdktrace.SpanExporter, error) {
 	// Just support jaeger and zipkin now, more for later
 	switch c.Batcher {
 	case kindJaeger:
-		if c.AgentHost != "" && c.AgentPort != "" {
+		if c.isAgentEndPoint() {
 			return jaeger.New(jaeger.WithAgentEndpoint(jaeger.WithAgentHost(c.AgentHost), jaeger.WithAgentPort(c.AgentPort)))
-		} else {
-			return jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.Endpoint)))
 		}
+		return jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.Endpoint)))
 	case kindZipkin:
 		return zipkin.New(c.Endpoint)
 	case kindOtlpGrpc:
diff --git a/core/trace/agent_test.go b/core/trace/agent_test.go
index 1d20694e..dc9172b4 100644
--- a/core/trace/agent_test.go
+++ b/core/trace/agent_test.go
@@ -11,10 +11,12 @@ func TestStartAgent(t *testing.T) {
 	logx.Disable()
 
 	const (
-		endpoint1 = "localhost:1234"
-		endpoint2 = "remotehost:1234"
-		endpoint3 = "localhost:1235"
-		endpoint4 = "localhost:1236"
+		endpoint1  = "localhost:1234"
+		endpoint2  = "remotehost:1234"
+		endpoint3  = "localhost:1235"
+		endpoint4  = "localhost:1236"
+		agentHost1 = "localhost"
+		agentPort1 = "6831"
 	)
 	c1 := Config{
 		Name: "foo",
@@ -44,6 +46,19 @@ func TestStartAgent(t *testing.T) {
 		Endpoint: endpoint4,
 		Batcher:  kindOtlpHttp,
 	}
+	c7 := Config{
+		Name:      "jaegerUDP",
+		AgentHost: agentHost1,
+		AgentPort: agentPort1,
+		Batcher:   kindJaeger,
+	}
+	c8 := Config{
+		Name:      "jaegerUDP",
+		AgentHost: agentHost1,
+		AgentPort: agentPort1,
+		Endpoint:  endpoint1,
+		Batcher:   kindJaeger,
+	}
 
 	StartAgent(c1)
 	StartAgent(c1)
@@ -52,16 +67,26 @@ func TestStartAgent(t *testing.T) {
 	StartAgent(c4)
 	StartAgent(c5)
 	StartAgent(c6)
+	StartAgent(c7)
+	StartAgent(c8)
 
 	lock.Lock()
 	defer lock.Unlock()
 
 	// because remotehost cannot be resolved
-	assert.Equal(t, 4, len(agents))
+	assert.Equal(t, 5, len(agents))
 	_, ok := agents[""]
 	assert.True(t, ok)
 	_, ok = agents[endpoint1]
 	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()]
+	assert.True(t, ok)
 }
diff --git a/core/trace/config.go b/core/trace/config.go
index 5346e469..4abf2e8b 100644
--- a/core/trace/config.go
+++ b/core/trace/config.go
@@ -1,5 +1,7 @@
 package trace
 
+import "fmt"
+
 // TraceName represents the tracing name.
 const TraceName = "go-zero"
 
@@ -10,5 +12,16 @@ type Config struct {
 	AgentPort string  `json:",optional"`
 	Endpoint  string  `json:",optional"`
 	Sampler   float64 `json:",default=1.0"`
-	Batcher   string  `json:",default=jaeger,options=jaeger|zipkin|otlpgrpc|otlphttp"`
+	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) getEndpoint() string {
+	if c.isAgentEndPoint() {
+		return fmt.Sprintf("%s:%s", c.AgentHost, c.AgentPort)
+	}
+	return c.Endpoint
 }