Create & Init Project...

app/service/main/dapper/server/udpcollect/BUILD

@@ -0,0 +1,38 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["udpcollect_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["udpcollect.go"],
+    importpath = "go-common/app/service/main/dapper/server/udpcollect",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = ["//library/log:go_default_library"],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

app/service/main/dapper/server/udpcollect/udpcollect.go

@@ -0,0 +1,156 @@
+package udpcollect
+
+import (
+	"fmt"
+	"net"
+	"net/url"
+	"os"
+	"path"
+	"strings"
+	"sync"
+	"time"
+
+	"go-common/library/log"
+)
+
+const (
+	_bufsize = 32 * 1024
+)
+
+// New UnixCollect
+func New(addr string, workers int, writeFn func(p []byte) error) (*UDPCollect, error) {
+	if workers == 0 {
+		workers = 1
+	}
+	addrURL, err := url.Parse(addr)
+	if err != nil {
+		return nil, fmt.Errorf("parse addr error: %s", err)
+	}
+	return &UDPCollect{
+		addr:    addrURL,
+		writeFn: writeFn,
+		workers: workers,
+		pool: sync.Pool{
+			New: func() interface{} {
+				return make([]byte, _bufsize)
+			},
+		},
+		readTimeout: 60 * time.Second,
+	}, nil
+}
+
+// UDPCollect collect span data from unix socket
+type UDPCollect struct {
+	wg          sync.WaitGroup
+	workers     int
+	addr        *url.URL
+	writeFn     func(p []byte) error
+	readTimeout time.Duration
+	pool        sync.Pool
+	closed      bool
+	pconn       net.PacketConn
+}
+
+// Start collector
+func (u *UDPCollect) Start() error {
+	var err error
+	switch u.addr.Scheme {
+	case "unixgram":
+		u.pconn, err = listenUNIX(u.addr.Path)
+	case "udp", "udp4", "udp6":
+		u.pconn, err = listtenNet(u.addr.Scheme, u.addr.Host)
+	default:
+		return fmt.Errorf("unsupport network %s", u.addr.Scheme)
+	}
+	if err != nil {
+		return fmt.Errorf("listen packet error: %s", err)
+	}
+	log.Info("dapper agent listen at: %s, workers: %d", u.addr, u.workers)
+	u.wg.Add(u.workers)
+	for i := 0; i < u.workers; i++ {
+		go u.serve()
+	}
+	return nil
+}
+
+func listenUNIX(addr string) (net.PacketConn, error) {
+	dirname := path.Dir(addr)
+	info, err := os.Stat(dirname)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			return nil, err
+		}
+		if err := os.MkdirAll(dirname, 0755); err != nil {
+			return nil, fmt.Errorf("create directory %s error: %s", dirname, err)
+		}
+	}
+	if err == nil && !info.IsDir() {
+		return nil, fmt.Errorf("%s is already exists and not a directory", dirname)
+	}
+	if _, err := os.Stat(addr); err == nil {
+		// remove old socket file
+		os.Remove(addr)
+	}
+	conn, err := net.ListenPacket("unixgram", addr)
+	if err != nil {
+		return nil, err
+	}
+	// make file permission to 666, so php can wirte span to this socket
+	return conn, os.Chmod(addr, 0666)
+}
+
+func listtenNet(network, addr string) (net.PacketConn, error) {
+	return net.ListenPacket(network, addr)
+}
+func (u *UDPCollect) serve() {
+	defer u.wg.Done()
+	for {
+		if err := u.handler(u.pconn); err != nil {
+			if strings.Contains(err.Error(), "closed") && u.closed {
+				return
+			}
+			log.Error("handler PacketConn error: %s, retry after second", err)
+			time.Sleep(time.Second)
+		}
+	}
+}
+
+func (u *UDPCollect) handler(pconn net.PacketConn) error {
+	p := u.buffer()
+	defer u.freeBuffer(p)
+	pconn.SetReadDeadline(time.Now().Add(u.readTimeout))
+	n, _, err := pconn.ReadFrom(p)
+	if n > 0 {
+		u.writeFn(p[:n])
+	}
+	if err == nil {
+		return nil
+	}
+	if netErr, ok := err.(net.Error); ok {
+		// ignore timeout and temporyary
+		if netErr.Timeout() || netErr.Temporary() {
+			return nil
+		}
+	}
+	return err
+}
+
+func (u *UDPCollect) buffer() []byte {
+	return u.pool.Get().([]byte)
+}
+
+func (u *UDPCollect) freeBuffer(p []byte) {
+	u.pool.Put(p)
+}
+
+// Close udp collect
+func (u *UDPCollect) Close() error {
+	u.closed = true
+	u.pconn.Close()
+	// wait all workers exit
+	u.wg.Wait()
+	if u.addr.Scheme == "unixgram" {
+		return os.Remove(u.addr.Path)
+	}
+	return nil
+}

app/service/main/dapper/server/udpcollect/udpcollect_test.go

@@ -0,0 +1,38 @@
+package udpcollect
+
+import (
+	"bytes"
+	"net"
+	"testing"
+	"time"
+)
+
+func TestUDPCollect(t *testing.T) {
+	count := 0
+	data := []byte("hello world")
+	collect, err := New("unixgram:///tmp/test.sock", 2, func(p []byte) error {
+		count++
+		if !bytes.Equal(p, data) {
+			t.Errorf("invalid p: %s", p)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := collect.Start(); err != nil {
+		t.Fatal(err)
+	}
+	conn, err := net.DialTimeout("unixgram", "/tmp/test.sock", time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 20; i++ {
+		conn.Write(data)
+	}
+	time.Sleep(time.Second)
+	collect.Close()
+	if count != 20 {
+		t.Errorf("wrong get %d != 20", count)
+	}
+}