Support connections over unix sockets

conn.go

@@ -9,7 +9,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/fs"
 	"net"
+	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -59,7 +61,7 @@ type Conn struct {
 
 	config *Config
 
-	conn    *net.TCPConn
+	conn    net.Conn
 	tlsConn *tls.Conn
 	addr    string
 
@@ -90,6 +92,7 @@ func NewConn(addr string, config *Config, delegate ConnDelegate) *Conn {
 	if !config.initialized {
 		panic("Config must be created with NewConfig()")
 	}
+
 	return &Conn{
 		addr: addr,
 
@@ -119,8 +122,7 @@ func NewConn(addr string, config *Config, delegate ConnDelegate) *Conn {
 // The logger parameter is an interface that requires the following
 // method to be implemented (such as the the stdlib log.Logger):
 //
-//    Output(calldepth int, s string)
-//
+//	Output(calldepth int, s string)
 func (c *Conn) SetLogger(l logger, lvl LogLevel, format string) {
 	c.logGuard.Lock()
 	defer c.logGuard.Unlock()
@@ -176,11 +178,11 @@ func (c *Conn) Connect() (*IdentifyResponse, error) {
 		Timeout:   c.config.DialTimeout,
 	}
 
-	conn, err := dialer.Dial("tcp", c.addr)
+	conn, err := dialer.Dial(c.socketType(), c.addr)
 	if err != nil {
 		return nil, err
 	}
-	c.conn = conn.(*net.TCPConn)
+	c.conn = conn
 	c.r = conn
 	c.w = conn
 
@@ -218,7 +220,7 @@ func (c *Conn) Connect() (*IdentifyResponse, error) {
 func (c *Conn) Close() error {
 	atomic.StoreInt32(&c.closeFlag, 1)
 	if c.conn != nil && atomic.LoadInt64(&c.messagesInFlight) == 0 {
-		return c.conn.CloseRead()
+		return c.conn.Close()
 	}
 	return nil
 }
@@ -415,6 +417,7 @@ func (c *Conn) identify() (*IdentifyResponse, error) {
 }
 
 func (c *Conn) upgradeTLS(tlsConf *tls.Config) error {
+	// TLS does not support unix sockets
 	host, _, err := net.SplitHostPort(c.addr)
 	if err != nil {
 		return err
@@ -426,7 +429,6 @@ func (c *Conn) upgradeTLS(tlsConf *tls.Config) error {
 		conf = tlsConf.Clone()
 	}
 	conf.ServerName = host
-
 	c.tlsConn = tls.Client(c.conn, conf)
 	err = c.tlsConn.Handshake()
 	if err != nil {
@@ -666,7 +668,7 @@ func (c *Conn) close() {
 	c.stopper.Do(func() {
 		c.log(LogLevelInfo, "beginning close")
 		close(c.exitChan)
-		c.conn.CloseRead()
+		c.conn.Close()
 
 		c.wg.Add(1)
 		go c.cleanup()
@@ -720,7 +722,7 @@ func (c *Conn) waitForCleanup() {
 	// this blocks until readLoop and writeLoop
 	// (and cleanup goroutine above) have exited
 	c.wg.Wait()
-	c.conn.CloseWrite()
+	c.conn.Close()
 	c.log(LogLevelInfo, "clean close complete")
 	c.delegate.OnClose(c)
 }
@@ -763,3 +765,18 @@ func (c *Conn) log(lvl LogLevel, line string, args ...interface{}) {
 		fmt.Sprintf(logFmt, c.String()),
 		fmt.Sprintf(line, args...)))
 }
+
+func (c *Conn) socketType() string {
+	if isSocket(c.addr) {
+		return "unix"
+	}
+	return "tcp"
+}
+
+func isSocket(path string) bool {
+	fileInfo, err := os.Stat(path)
+	if err != nil {
+		return false
+	}
+	return fileInfo.Mode().Type() == fs.ModeSocket
+}

consumer_unix_socket_test.go

@@ -0,0 +1,140 @@
+package nsq
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"net"
+	"net/http"
+	"strconv"
+	"testing"
+	"time"
+)
+
+func SendMessageToUnixSocket(t *testing.T, addr string, topic string, method string, body []byte) {
+	httpclient := http.Client{
+		Transport: &http.Transport{
+			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+				return net.Dial("unix", addr)
+			},
+		},
+	}
+
+	endpoint := fmt.Sprintf("http://unix/%s?topic=%s", method, topic)
+	resp, err := httpclient.Post(endpoint, "application/octet-stream", bytes.NewBuffer(body))
+	if err != nil {
+		t.Fatalf(err.Error())
+		return
+	}
+	if resp.StatusCode != 200 {
+		t.Fatalf("%s status code: %d", method, resp.StatusCode)
+	}
+	resp.Body.Close()
+}
+
+func TestUnixSocketConsumer(t *testing.T) {
+	consumerUnixSocketTest(t, nil)
+}
+
+func TestUnixSocketConsumerDeflate(t *testing.T) {
+	consumerUnixSocketTest(t, func(c *Config) {
+		c.Deflate = true
+	})
+}
+
+func TestUnixSocketConsumerSnappy(t *testing.T) {
+	consumerUnixSocketTest(t, func(c *Config) {
+		c.Snappy = true
+	})
+}
+
+func consumerUnixSocketTest(t *testing.T, cb func(c *Config)) {
+	// unix addresses of nsqd are specified in test.sh
+	addr := "/tmp/nsqd.sock"
+	httpAddr := "/tmp/nsqd-http.sock"
+
+	config := NewConfig()
+
+	config.DefaultRequeueDelay = 0
+	// so that the test won't timeout from backing off
+	config.MaxBackoffDuration = time.Millisecond * 50
+	if cb != nil {
+		cb(config)
+	}
+	topicName := "rdr_test"
+	if config.Deflate {
+		topicName = topicName + "_deflate"
+	} else if config.Snappy {
+		topicName = topicName + "_snappy"
+	}
+	topicName = topicName + strconv.Itoa(int(time.Now().Unix()))
+	q, _ := NewConsumer(topicName, "ch", config)
+	q.SetLogger(newTestLogger(t), LogLevelDebug)
+
+	h := &MyTestHandler{
+		t: t,
+		q: q,
+	}
+	q.AddHandler(h)
+
+	SendMessageToUnixSocket(t, httpAddr, topicName, "pub", []byte(`{"msg":"single"}`))
+	SendMessageToUnixSocket(t, httpAddr, topicName, "mpub", []byte("{\"msg\":\"double\"}\n{\"msg\":\"double\"}"))
+	SendMessageToUnixSocket(t, httpAddr, topicName, "pub", []byte("TOBEFAILED"))
+	h.messagesSent = 4
+
+	err := q.ConnectToNSQD(addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	stats := q.Stats()
+	if stats.Connections == 0 {
+		t.Fatal("stats report 0 connections (should be > 0)")
+	}
+
+	err = q.ConnectToNSQD(addr)
+	if err == nil {
+		t.Fatal("should not be able to connect to the same NSQ twice")
+	}
+
+	conn := q.conns()[0]
+	if conn.String() != addr {
+		t.Fatal("connection should be bound to the specified address:", addr)
+	}
+
+	err = q.DisconnectFromNSQD("/tmp/doesntexist.sock")
+	if err == nil {
+		t.Fatal("should not be able to disconnect from an unknown nsqd")
+	}
+
+	err = q.ConnectToNSQD("/tmp/doesntexist.sock")
+	if err == nil {
+		t.Fatal("should not be able to connect to non-existent nsqd")
+	}
+
+	err = q.DisconnectFromNSQD("/tmp/doesntexist.sock")
+	if err != nil {
+		t.Fatal("should be able to disconnect from an nsqd - " + err.Error())
+	}
+
+	<-q.StopChan
+
+	stats = q.Stats()
+	if stats.Connections != 0 {
+		t.Fatalf("stats report %d active connections (should be 0)", stats.Connections)
+	}
+
+	stats = q.Stats()
+	if stats.MessagesReceived != uint64(h.messagesReceived+h.messagesFailed) {
+		t.Fatalf("stats report %d messages received (should be %d)",
+			stats.MessagesReceived,
+			h.messagesReceived+h.messagesFailed)
+	}
+
+	if h.messagesReceived != 8 || h.messagesSent != 4 {
+		t.Fatalf("end of test. should have handled a diff number of messages (got %d, sent %d)", h.messagesReceived, h.messagesSent)
+	}
+	if h.messagesFailed != 1 {
+		t.Fatal("failed message not done")
+	}
+}