Repository: atlassian/gostatsd
Pull request #472 (2022-03-18 10:04:33)

Allow using Unix Domain Socket instead of UDP
Allow using Unix Domain Sockets instead of UDP for the server.
If the `metrics-addr` is not empty and not `[host]:port` combination it will be considered a Unix Socket path and it will create and listen in that socket instead of UDP.

atlassian-cla-bot[bot]: Thank you for your submission! Like many open source projects, we ask that you sign our <a href=https://opensource.atlassian.com/cla>CLA (Contributor License Agreement)</a> before we can accept your contribution.
      <strong>If your email is listed below</strong>, please ensure that you sign the CLA with the <strong>same email address</strong>.</br>
      <strong>The following users still need to sign our CLA:</strong></br>❌rubenruizdegauna
<p><sub>Already signed the CLA? To re-check, try refreshing the page.</sub></p>
rubenruizdegauna: Thanks @aelse @tiedotguy for the quick review. I have made the requested changes in https://github.com/atlassian/gostatsd/pull/472/commits/11d015b2bd6251deb18b8db3ae611374181465aa
aelse: lgtm. If @tiedotguy is happy we can merge.
rubenruizdegauna: > lgtm. If @tiedotguy is happy we can merge.

Hi @tiedotguy , do you think this can be merged? Thanks!

Patch:
diff --git a/README.md b/README.md
index a9dd72c..108d5af 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,10 @@ to run `gostatsd` with a graphite backend and a grafana dashboard.
 While not generally tested on Windows, it should work.  Maximum throughput is likely to be better on
 a linux system, however.
 
+The server listens for UDP packets by default. You can use unix sockets providing an absolute path to the socket 
+in the `metrics-addr` configuration option. The socket mode used in this case is SOCK_DGRAM.
+Note that using unix sockets will only work on linux and that it will ignore `conn-per-reader` configuration option.
+
 Configuring the server mode
 ---------------------------
 The server can currently run in two modes: `standalone` and `forwarder`.  It is configured through the top level
@@ -83,7 +87,8 @@ This configuration mode allows the following configuration options:
   so that memory can be pre-allocated and reducing churn.  Defaults to `4`.  Note: this is only a hint, and it is safe
   to send more.
 - `log-raw-metric`: logs raw metrics received from the network.  Defaults to `false`.
-- `metrics-addr`: the address to listen to metrics on. Defaults to `:8125`.
+- `metrics-addr`: the address to listen to metrics on. Defaults to `:8125`. Using a file path instead of `host:port` 
+  will create a Unix Domain Socket in the specified path instead of using UDP.
 - `namespace`: a namespace to prefix all metrics with.  Defaults to ''.
 - `statser-type`: configures where internal metrics are sent to.  May be `internal` which sends them to the internal
   processing pipeline, `logging` which logs them, `null` which drops them.  Defaults to `internal`, or `null` if the
@@ -93,7 +98,8 @@ This configuration mode allows the following configuration options:
   Defaults to `false`.
 - `receive-batch-size`: the number of datagrams to attempt to read.  It is more CPU efficient to read multiple, however
   it takes extra memory.  See [Memory allocation for read buffers] section below for details.  Defaults to 50.
-- `conn-per-reader`: attempts to create a connection for every UDP receiver.  Not supported by all OS versions.
+- `conn-per-reader`: attempts to create a connection for every UDP receiver.  Not supported by all OS versions. It will 
+  be ignored when unix sockets are used for the connection.
   Defaults to `false`.
 - `bad-lines-per-minute`: the number of metrics which fail to parse to log per minute.  This is used to prevent a bad
   client spamming malformed statsd data, while still logging some information to enable troubleshooting.  Defaults to `0`.
diff --git a/pkg/statsd/receiver.go b/pkg/statsd/receiver.go
index 700d07b..ad3bc6a 100644
--- a/pkg/statsd/receiver.go
+++ b/pkg/statsd/receiver.go
@@ -3,6 +3,7 @@ package statsd
 import (
 	"context"
 	"net"
+	"os"
 	"strings"
 	"sync/atomic"
 
@@ -98,6 +99,13 @@ func (dr *DatagramReceiver) Run(ctx context.Context) {
 		}
 	}
 
+	// Delete socket file when connection is made using Unix Domain Sockets
+	if len(connections) > 0 {
+		if socket, ok := connections[0].LocalAddr().(*net.UnixAddr); ok {
+			defer os.Remove(socket.String())
+		}
+	}
+
 	// Wait for everything to stop
 	wg.Wait()
 }
@@ -142,8 +150,14 @@ func (dr *DatagramReceiver) Receive(ctx context.Context, c net.PacketConn) {
 				dr.bufPool.Put(retBuf)
 			}
 
+			ip := gostatsd.UnknownSource
+			// Do not retrieve IP address when connection is made using Unix Domain Sockets
+			if _, isUnixAddr := c.LocalAddr().(*net.UnixAddr); !isUnixAddr {
+				ip = getIP(addr)
+			}
+
 			dgs[i] = &Datagram{
-				IP:        getIP(addr),
+				IP:        ip,
 				Msg:       buf,
 				Timestamp: now,
 				DoneFunc:  doneFn,
diff --git a/pkg/statsd/statsd.go b/pkg/statsd/statsd.go
index 8390571..d9c5093 100644
--- a/pkg/statsd/statsd.go
+++ b/pkg/statsd/statsd.go
@@ -81,12 +81,22 @@ func socketFactory(metricsAddr string, connPerReader bool) SocketFactory {
 		return func() (net.PacketConn, error) {
 			return reuseport.ListenPacket("udp", metricsAddr)
 		}
-	} else {
-		conn, err := net.ListenPacket("udp", metricsAddr)
-		return func() (net.PacketConn, error) {
-			return conn, err
-		}
 	}
+
+	conn, err := net.ListenPacket(networkFromAddress(metricsAddr), metricsAddr)
+	return func() (net.PacketConn, error) {
+		return conn, err
+	}
+}
+
+//networkFromAddress returns the network type based on the provided address
+//if the address starts with a slash (unix absolute path) it will be considered a unix socket
+//otherwise it will default to UDP
+func networkFromAddress(addr string) string {
+	if len(addr) > 0 && addr[0:1] == "/" {
+		return "unixgram"
+	}
+	return "udp"
 }
 
 func (s *Server) createStandaloneSink() (gostatsd.PipelineHandler, []gostatsd.Runnable, error) {


Test patch:
diff --git a/pkg/statsd/receiver_test.go b/pkg/statsd/receiver_test.go
index 830bfbc..ebb6e0c 100644
--- a/pkg/statsd/receiver_test.go
+++ b/pkg/statsd/receiver_test.go
@@ -2,12 +2,15 @@ package statsd
 
 import (
 	"context"
+	"net"
+	"os"
 	"runtime"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/magiconair/properties/assert"
+	tassert "github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/atlassian/gostatsd"
@@ -93,3 +96,72 @@ func TestDatagramReceiver_Receive(t *testing.T) {
 	assert.Equal(t, string(dg.IP), fakesocket.FakeAddr.IP.String())
 	assert.Equal(t, dg.Msg, fakesocket.FakeMetric)
 }
+
+func TestDatagramReceiver_UnixSocketConnection(t *testing.T) {
+	ch := make(chan []*Datagram, 1)
+	message := "abc.def.g:10|c"
+
+	// Datagram receiver listening in Unix Domain Socket
+	socketPath := os.TempDir() + "/gostatsd_receiver_test_receive_uds.sock"
+	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		mr.Run(ctx)
+		wg.Done()
+	}()
+
+	// Wait until the socket is created
+	tassert.Eventually(t, func() bool {
+		return tassert.FileExists(t, socketPath)
+	}, time.Second, 10*time.Millisecond)
+
+	err := sendDataToSocket(socketPath, message)
+	tassert.NoError(t, err)
+
+	select {
+	case d := <-ch:
+		tassert.Len(t, d, 1)
+		tassert.Equal(t, d[0].Msg, []byte(message))
+		cancel()
+	case <-time.After(time.Second):
+		t.Errorf("Timeout, failed to read datagram")
+		cancel()
+	}
+	wg.Wait()
+}
+
+func TestDatagramReceiver_UnixSocketIsRemovedOnContextCancellation(t *testing.T) {
+	ch := make(chan []*Datagram, 1)
+
+	socketPath := os.TempDir() + "/gostatsd_receiver_test_receive_uds.sock"
+	mr := NewDatagramReceiver(ch, socketFactory(socketPath, false), 1, 2)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		mr.Run(ctx)
+		wg.Done()
+	}()
+
+	tassert.Eventually(t, func() bool {
+		return tassert.FileExists(t, socketPath)
+	}, time.Second, 10*time.Millisecond)
+	cancel()
+	wg.Wait()
+	tassert.NoFileExists(t, socketPath)
+}
+
+func sendDataToSocket(socketPath string, data string) error {
+	c, err := net.Dial("unixgram", socketPath)
+	if err != nil {
+		return err
+	}
+	defer c.Close()
+
+	_, err = c.Write([]byte(data))
+	return err
+}
diff --git a/pkg/statsd/statsd_test.go b/pkg/statsd/statsd_test.go
index 2441a44..c05f8ae 100644
--- a/pkg/statsd/statsd_test.go
+++ b/pkg/statsd/statsd_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"math/rand"
 	"runtime"
+	"strconv"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -11,6 +12,7 @@ import (
 	"github.com/ash2k/stager/wait"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/time/rate"
 
@@ -145,6 +147,35 @@ func TestStatsdThroughput(t *testing.T) {
 		memStatsFinish.GCCPUFraction)
 }
 
+func TestNetworkFromAddress(t *testing.T) {
+	t.Parallel()
+	input := []struct {
+		address         string
+		expectedNetwork string
+	}{
+		{
+			address:         "localhost:8125",
+			expectedNetwork: "udp",
+		},
+		{
+			address:         "",
+			expectedNetwork: "udp",
+		},
+		{
+			address:         "/some/file.sock",
+			expectedNetwork: "unixgram",
+		},
+	}
+	for pos, inp := range input {
+		inp := inp
+		t.Run(strconv.Itoa(pos), func(t *testing.T) {
+			t.Parallel()
+			network := networkFromAddress(inp.address)
+			assert.Equal(t, inp.expectedNetwork, network)
+		})
+	}
+}
+
 type countingBackend struct {
 	metrics uint64
 	events  uint64