github.com/nsqio/nsq@v1.3.0/nsqd/protocol_v2_unixsocket_test.go (about)

     1  package nsqd
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"compress/flate"
     7  	"crypto/tls"
     8  	"encoding/json"
     9  	"errors"
    10  	"fmt"
    11  	"math"
    12  	"math/rand"
    13  	"net"
    14  	"os"
    15  	"path"
    16  	"runtime"
    17  	"strconv"
    18  	"sync"
    19  	"sync/atomic"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/golang/snappy"
    24  	"github.com/nsqio/go-nsq"
    25  	"github.com/nsqio/nsq/internal/protocol"
    26  	"github.com/nsqio/nsq/internal/test"
    27  )
    28  
    29  func mustUnixSocketStartNSQD(opts *Options) (net.Addr, net.Addr, *NSQD) {
    30  	tmpDir := os.TempDir()
    31  	opts.TCPAddress = path.Join(tmpDir, fmt.Sprintf("nsqd-%d.sock", rand.Int()))
    32  	opts.HTTPAddress = path.Join(tmpDir, fmt.Sprintf("nsqd-%d.sock", rand.Int()))
    33  
    34  	if opts.DataPath == "" {
    35  		tmpDir, err := os.MkdirTemp("", "nsq-test-")
    36  		if err != nil {
    37  			panic(err)
    38  		}
    39  		opts.DataPath = tmpDir
    40  	}
    41  	nsqd, err := New(opts)
    42  	if err != nil {
    43  		panic(err)
    44  	}
    45  	go func() {
    46  		err := nsqd.Main()
    47  		if err != nil {
    48  			panic(err)
    49  		}
    50  	}()
    51  	return nsqd.RealTCPAddr(), nsqd.RealHTTPAddr(), nsqd
    52  }
    53  
    54  func mustUnixSocketConnectNSQD(addr net.Addr) (net.Conn, error) {
    55  	conn, err := net.DialTimeout("unix", addr.String(), time.Second)
    56  	if err != nil {
    57  		return nil, err
    58  	}
    59  	conn.Write(nsq.MagicV2)
    60  	return conn, nil
    61  }
    62  
    63  // exercise the basic operations of the V2 protocol
    64  func TestUnixSocketBasicV2(t *testing.T) {
    65  	opts := NewOptions()
    66  	opts.Logger = test.NewTestLogger(t)
    67  	opts.ClientTimeout = 60 * time.Second
    68  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
    69  	defer os.RemoveAll(opts.DataPath)
    70  	defer nsqd.Exit()
    71  
    72  	topicName := "test_v2" + strconv.Itoa(int(time.Now().Unix()))
    73  	topic := nsqd.GetTopic(topicName)
    74  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
    75  	topic.PutMessage(msg)
    76  
    77  	conn, err := mustUnixSocketConnectNSQD(addr)
    78  	test.Nil(t, err)
    79  	defer conn.Close()
    80  
    81  	identify(t, conn, nil, frameTypeResponse)
    82  	sub(t, conn, topicName, "ch")
    83  
    84  	_, err = nsq.Ready(1).WriteTo(conn)
    85  	test.Nil(t, err)
    86  
    87  	resp, err := nsq.ReadResponse(conn)
    88  	test.Nil(t, err)
    89  	frameType, data, _ := nsq.UnpackResponse(resp)
    90  	msgOut, _ := decodeMessage(data)
    91  	test.Equal(t, frameTypeMessage, frameType)
    92  	test.Equal(t, msg.ID, msgOut.ID)
    93  	test.Equal(t, msg.Body, msgOut.Body)
    94  	test.Equal(t, uint16(1), msgOut.Attempts)
    95  }
    96  
    97  func TestUnixSocketMultipleConsumerV2(t *testing.T) {
    98  	msgChan := make(chan *Message)
    99  
   100  	opts := NewOptions()
   101  	opts.Logger = test.NewTestLogger(t)
   102  	opts.ClientTimeout = 60 * time.Second
   103  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   104  	defer os.RemoveAll(opts.DataPath)
   105  	defer nsqd.Exit()
   106  
   107  	topicName := "test_multiple_v2" + strconv.Itoa(int(time.Now().Unix()))
   108  	topic := nsqd.GetTopic(topicName)
   109  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
   110  	topic.GetChannel("ch1")
   111  	topic.GetChannel("ch2")
   112  	topic.PutMessage(msg)
   113  
   114  	for _, i := range []string{"1", "2"} {
   115  		conn, err := mustUnixSocketConnectNSQD(addr)
   116  		test.Nil(t, err)
   117  		defer conn.Close()
   118  
   119  		identify(t, conn, nil, frameTypeResponse)
   120  		sub(t, conn, topicName, "ch"+i)
   121  
   122  		_, err = nsq.Ready(1).WriteTo(conn)
   123  		test.Nil(t, err)
   124  
   125  		go func(c net.Conn) {
   126  			resp, err := nsq.ReadResponse(c)
   127  			test.Nil(t, err)
   128  			_, data, err := nsq.UnpackResponse(resp)
   129  			test.Nil(t, err)
   130  			msg, err := decodeMessage(data)
   131  			test.Nil(t, err)
   132  			msgChan <- msg
   133  		}(conn)
   134  	}
   135  
   136  	msgOut := <-msgChan
   137  	test.Equal(t, msg.ID, msgOut.ID)
   138  	test.Equal(t, msg.Body, msgOut.Body)
   139  	test.Equal(t, uint16(1), msgOut.Attempts)
   140  	msgOut = <-msgChan
   141  	test.Equal(t, msg.ID, msgOut.ID)
   142  	test.Equal(t, msg.Body, msgOut.Body)
   143  	test.Equal(t, uint16(1), msgOut.Attempts)
   144  }
   145  
   146  func TestUnixSocketClientTimeout(t *testing.T) {
   147  	topicName := "test_client_timeout_v2" + strconv.Itoa(int(time.Now().Unix()))
   148  
   149  	opts := NewOptions()
   150  	opts.Logger = test.NewTestLogger(t)
   151  	opts.ClientTimeout = 150 * time.Millisecond
   152  	opts.LogLevel = LOG_DEBUG
   153  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   154  	defer os.RemoveAll(opts.DataPath)
   155  	defer nsqd.Exit()
   156  
   157  	conn, err := mustUnixSocketConnectNSQD(addr)
   158  	test.Nil(t, err)
   159  	defer conn.Close()
   160  
   161  	identify(t, conn, nil, frameTypeResponse)
   162  	sub(t, conn, topicName, "ch")
   163  
   164  	time.Sleep(150 * time.Millisecond)
   165  
   166  	// depending on timing there may be 1 or 2 hearbeats sent
   167  	// just read until we get an error
   168  	timer := time.After(100 * time.Millisecond)
   169  	for {
   170  		select {
   171  		case <-timer:
   172  			t.Fatalf("test timed out")
   173  		default:
   174  			_, err := nsq.ReadResponse(conn)
   175  			if err != nil {
   176  				goto done
   177  			}
   178  		}
   179  	}
   180  done:
   181  }
   182  
   183  func TestUnixSocketClientHeartbeat(t *testing.T) {
   184  	topicName := "test_hb_v2" + strconv.Itoa(int(time.Now().Unix()))
   185  
   186  	opts := NewOptions()
   187  	opts.Logger = test.NewTestLogger(t)
   188  	opts.ClientTimeout = 200 * time.Millisecond
   189  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   190  	defer os.RemoveAll(opts.DataPath)
   191  	defer nsqd.Exit()
   192  
   193  	conn, err := mustUnixSocketConnectNSQD(addr)
   194  	test.Nil(t, err)
   195  	defer conn.Close()
   196  
   197  	identify(t, conn, nil, frameTypeResponse)
   198  	sub(t, conn, topicName, "ch")
   199  
   200  	_, err = nsq.Ready(1).WriteTo(conn)
   201  	test.Nil(t, err)
   202  
   203  	resp, _ := nsq.ReadResponse(conn)
   204  	_, data, _ := nsq.UnpackResponse(resp)
   205  	test.Equal(t, []byte("_heartbeat_"), data)
   206  
   207  	time.Sleep(20 * time.Millisecond)
   208  
   209  	_, err = nsq.Nop().WriteTo(conn)
   210  	test.Nil(t, err)
   211  
   212  	// wait long enough that would have timed out (had we not sent the above cmd)
   213  	time.Sleep(100 * time.Millisecond)
   214  
   215  	_, err = nsq.Nop().WriteTo(conn)
   216  	test.Nil(t, err)
   217  }
   218  
   219  func TestUnixSocketClientHeartbeatDisableSUB(t *testing.T) {
   220  	topicName := "test_hb_v2" + strconv.Itoa(int(time.Now().Unix()))
   221  
   222  	opts := NewOptions()
   223  	opts.Logger = test.NewTestLogger(t)
   224  	opts.ClientTimeout = 200 * time.Millisecond
   225  	opts.LogLevel = LOG_DEBUG
   226  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   227  	defer os.RemoveAll(opts.DataPath)
   228  	defer nsqd.Exit()
   229  
   230  	conn, err := mustUnixSocketConnectNSQD(addr)
   231  	test.Nil(t, err)
   232  	defer conn.Close()
   233  
   234  	identify(t, conn, map[string]interface{}{
   235  		"heartbeat_interval": -1,
   236  	}, frameTypeResponse)
   237  	subFail(t, conn, topicName, "ch")
   238  }
   239  
   240  func TestUnixSocketClientHeartbeatDisable(t *testing.T) {
   241  	opts := NewOptions()
   242  	opts.Logger = test.NewTestLogger(t)
   243  	opts.ClientTimeout = 100 * time.Millisecond
   244  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   245  	defer os.RemoveAll(opts.DataPath)
   246  	defer nsqd.Exit()
   247  
   248  	conn, err := mustUnixSocketConnectNSQD(addr)
   249  	test.Nil(t, err)
   250  	defer conn.Close()
   251  
   252  	identify(t, conn, map[string]interface{}{
   253  		"heartbeat_interval": -1,
   254  	}, frameTypeResponse)
   255  
   256  	time.Sleep(150 * time.Millisecond)
   257  
   258  	_, err = nsq.Nop().WriteTo(conn)
   259  	test.Nil(t, err)
   260  }
   261  
   262  func TestUnixSocketMaxHeartbeatIntervalValid(t *testing.T) {
   263  	opts := NewOptions()
   264  	opts.Logger = test.NewTestLogger(t)
   265  	opts.MaxHeartbeatInterval = 300 * time.Second
   266  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   267  	defer os.RemoveAll(opts.DataPath)
   268  	defer nsqd.Exit()
   269  
   270  	conn, err := mustUnixSocketConnectNSQD(addr)
   271  	test.Nil(t, err)
   272  	defer conn.Close()
   273  
   274  	hbi := int(opts.MaxHeartbeatInterval / time.Millisecond)
   275  	identify(t, conn, map[string]interface{}{
   276  		"heartbeat_interval": hbi,
   277  	}, frameTypeResponse)
   278  }
   279  
   280  func TestUnixSocketMaxHeartbeatIntervalInvalid(t *testing.T) {
   281  	opts := NewOptions()
   282  	opts.Logger = test.NewTestLogger(t)
   283  	opts.MaxHeartbeatInterval = 300 * time.Second
   284  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   285  	defer os.RemoveAll(opts.DataPath)
   286  	defer nsqd.Exit()
   287  
   288  	conn, err := mustUnixSocketConnectNSQD(addr)
   289  	test.Nil(t, err)
   290  	defer conn.Close()
   291  
   292  	hbi := int(opts.MaxHeartbeatInterval/time.Millisecond + 1)
   293  	data := identify(t, conn, map[string]interface{}{
   294  		"heartbeat_interval": hbi,
   295  	}, frameTypeError)
   296  	test.Equal(t, "E_BAD_BODY IDENTIFY heartbeat interval (300001) is invalid", string(data))
   297  }
   298  
   299  func TestUnixSocketPausing(t *testing.T) {
   300  	topicName := "test_pause_v2" + strconv.Itoa(int(time.Now().Unix()))
   301  
   302  	opts := NewOptions()
   303  	opts.Logger = test.NewTestLogger(t)
   304  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   305  	defer os.RemoveAll(opts.DataPath)
   306  	defer nsqd.Exit()
   307  
   308  	conn, err := mustUnixSocketConnectNSQD(addr)
   309  	test.Nil(t, err)
   310  	defer conn.Close()
   311  
   312  	identify(t, conn, nil, frameTypeResponse)
   313  	sub(t, conn, topicName, "ch")
   314  
   315  	_, err = nsq.Ready(1).WriteTo(conn)
   316  	test.Nil(t, err)
   317  
   318  	topic := nsqd.GetTopic(topicName)
   319  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
   320  	channel := topic.GetChannel("ch")
   321  	topic.PutMessage(msg)
   322  
   323  	// receive the first message via the client, finish it, and send new RDY
   324  	resp, _ := nsq.ReadResponse(conn)
   325  	_, data, _ := nsq.UnpackResponse(resp)
   326  	msg, _ = decodeMessage(data)
   327  	test.Equal(t, []byte("test body"), msg.Body)
   328  
   329  	_, err = nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(conn)
   330  	test.Nil(t, err)
   331  
   332  	_, err = nsq.Ready(1).WriteTo(conn)
   333  	test.Nil(t, err)
   334  
   335  	// sleep to allow the RDY state to take effect
   336  	time.Sleep(50 * time.Millisecond)
   337  
   338  	// pause the channel... the client shouldn't receive any more messages
   339  	channel.Pause()
   340  
   341  	// sleep to allow the paused state to take effect
   342  	time.Sleep(50 * time.Millisecond)
   343  
   344  	msg = NewMessage(topic.GenerateID(), []byte("test body2"))
   345  	topic.PutMessage(msg)
   346  
   347  	// allow the client to possibly get a message, the test would hang indefinitely
   348  	// if pausing was not working
   349  	time.Sleep(50 * time.Millisecond)
   350  	msg = <-channel.memoryMsgChan
   351  	test.Equal(t, []byte("test body2"), msg.Body)
   352  
   353  	// unpause the channel... the client should now be pushed a message
   354  	channel.UnPause()
   355  
   356  	msg = NewMessage(topic.GenerateID(), []byte("test body3"))
   357  	topic.PutMessage(msg)
   358  
   359  	resp, _ = nsq.ReadResponse(conn)
   360  	_, data, _ = nsq.UnpackResponse(resp)
   361  	msg, _ = decodeMessage(data)
   362  	test.Equal(t, []byte("test body3"), msg.Body)
   363  }
   364  
   365  func TestUnixSocketEmptyCommand(t *testing.T) {
   366  	opts := NewOptions()
   367  	opts.Logger = test.NewTestLogger(t)
   368  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   369  	defer os.RemoveAll(opts.DataPath)
   370  	defer nsqd.Exit()
   371  
   372  	conn, err := mustUnixSocketConnectNSQD(addr)
   373  	test.Nil(t, err)
   374  	defer conn.Close()
   375  
   376  	_, err = conn.Write([]byte("\n\n"))
   377  	test.Nil(t, err)
   378  
   379  	// if we didn't panic here we're good, see issue #120
   380  }
   381  
   382  func TestUnixSocketSizeLimits(t *testing.T) {
   383  	opts := NewOptions()
   384  	opts.Logger = test.NewTestLogger(t)
   385  	opts.LogLevel = LOG_DEBUG
   386  	opts.MaxMsgSize = 100
   387  	opts.MaxBodySize = 1000
   388  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   389  	defer os.RemoveAll(opts.DataPath)
   390  	defer nsqd.Exit()
   391  
   392  	conn, err := mustUnixSocketConnectNSQD(addr)
   393  	test.Nil(t, err)
   394  	defer conn.Close()
   395  
   396  	topicName := "test_limits_v2" + strconv.Itoa(int(time.Now().Unix()))
   397  
   398  	identify(t, conn, nil, frameTypeResponse)
   399  	sub(t, conn, topicName, "ch")
   400  
   401  	// PUB that's valid
   402  	nsq.Publish(topicName, make([]byte, 95)).WriteTo(conn)
   403  	resp, _ := nsq.ReadResponse(conn)
   404  	frameType, data, _ := nsq.UnpackResponse(resp)
   405  	t.Logf("frameType: %d, data: %s", frameType, data)
   406  	test.Equal(t, frameTypeResponse, frameType)
   407  	test.Equal(t, []byte("OK"), data)
   408  
   409  	// PUB that's invalid (too big)
   410  	nsq.Publish(topicName, make([]byte, 105)).WriteTo(conn)
   411  	resp, _ = nsq.ReadResponse(conn)
   412  	frameType, data, _ = nsq.UnpackResponse(resp)
   413  	t.Logf("frameType: %d, data: %s", frameType, data)
   414  	test.Equal(t, frameTypeError, frameType)
   415  	test.Equal(t, "E_BAD_MESSAGE PUB message too big 105 > 100", string(data))
   416  
   417  	// need to reconnect
   418  	conn, err = mustUnixSocketConnectNSQD(addr)
   419  	test.Nil(t, err)
   420  	defer conn.Close()
   421  
   422  	// PUB thats empty
   423  	nsq.Publish(topicName, []byte{}).WriteTo(conn)
   424  	resp, _ = nsq.ReadResponse(conn)
   425  	frameType, data, _ = nsq.UnpackResponse(resp)
   426  	t.Logf("frameType: %d, data: %s", frameType, data)
   427  	test.Equal(t, frameTypeError, frameType)
   428  	test.Equal(t, "E_BAD_MESSAGE PUB invalid message body size 0", string(data))
   429  
   430  	// need to reconnect
   431  	conn, err = mustUnixSocketConnectNSQD(addr)
   432  	test.Nil(t, err)
   433  	defer conn.Close()
   434  
   435  	// MPUB body that's valid
   436  	mpub := make([][]byte, 5)
   437  	for i := range mpub {
   438  		mpub[i] = make([]byte, 100)
   439  	}
   440  	cmd, _ := nsq.MultiPublish(topicName, mpub)
   441  	cmd.WriteTo(conn)
   442  	resp, _ = nsq.ReadResponse(conn)
   443  	frameType, data, _ = nsq.UnpackResponse(resp)
   444  	t.Logf("frameType: %d, data: %s", frameType, data)
   445  	test.Equal(t, frameTypeResponse, frameType)
   446  	test.Equal(t, []byte("OK"), data)
   447  
   448  	// MPUB body that's invalid (body too big)
   449  	mpub = make([][]byte, 11)
   450  	for i := range mpub {
   451  		mpub[i] = make([]byte, 100)
   452  	}
   453  	cmd, _ = nsq.MultiPublish(topicName, mpub)
   454  	cmd.WriteTo(conn)
   455  	resp, _ = nsq.ReadResponse(conn)
   456  	frameType, data, _ = nsq.UnpackResponse(resp)
   457  	t.Logf("frameType: %d, data: %s", frameType, data)
   458  	test.Equal(t, frameTypeError, frameType)
   459  	test.Equal(t, "E_BAD_BODY MPUB body too big 1148 > 1000", string(data))
   460  
   461  	// need to reconnect
   462  	conn, err = mustUnixSocketConnectNSQD(addr)
   463  	test.Nil(t, err)
   464  	defer conn.Close()
   465  
   466  	// MPUB that's invalid (one message empty)
   467  	mpub = make([][]byte, 5)
   468  	for i := range mpub {
   469  		mpub[i] = make([]byte, 100)
   470  	}
   471  	mpub = append(mpub, []byte{})
   472  	cmd, _ = nsq.MultiPublish(topicName, mpub)
   473  	cmd.WriteTo(conn)
   474  	resp, _ = nsq.ReadResponse(conn)
   475  	frameType, data, _ = nsq.UnpackResponse(resp)
   476  	t.Logf("frameType: %d, data: %s", frameType, data)
   477  	test.Equal(t, frameTypeError, frameType)
   478  	test.Equal(t, "E_BAD_MESSAGE MPUB invalid message(5) body size 0", string(data))
   479  
   480  	// need to reconnect
   481  	conn, err = mustUnixSocketConnectNSQD(addr)
   482  	test.Nil(t, err)
   483  	defer conn.Close()
   484  
   485  	// MPUB body that's invalid (one of the messages is too big)
   486  	mpub = make([][]byte, 5)
   487  	for i := range mpub {
   488  		mpub[i] = make([]byte, 101)
   489  	}
   490  	cmd, _ = nsq.MultiPublish(topicName, mpub)
   491  	cmd.WriteTo(conn)
   492  	resp, _ = nsq.ReadResponse(conn)
   493  	frameType, data, _ = nsq.UnpackResponse(resp)
   494  	t.Logf("frameType: %d, data: %s", frameType, data)
   495  	test.Equal(t, frameTypeError, frameType)
   496  	test.Equal(t, "E_BAD_MESSAGE MPUB message too big 101 > 100", string(data))
   497  }
   498  
   499  func TestUnixSocketDPUB(t *testing.T) {
   500  	opts := NewOptions()
   501  	opts.Logger = test.NewTestLogger(t)
   502  	opts.LogLevel = LOG_DEBUG
   503  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   504  	defer os.RemoveAll(opts.DataPath)
   505  	defer nsqd.Exit()
   506  
   507  	conn, err := mustUnixSocketConnectNSQD(addr)
   508  	test.Nil(t, err)
   509  	defer conn.Close()
   510  
   511  	topicName := "test_dpub_v2" + strconv.Itoa(int(time.Now().Unix()))
   512  
   513  	identify(t, conn, nil, frameTypeResponse)
   514  	sub(t, conn, topicName, "ch")
   515  
   516  	// valid
   517  	nsq.DeferredPublish(topicName, time.Second, make([]byte, 100)).WriteTo(conn)
   518  	resp, _ := nsq.ReadResponse(conn)
   519  	frameType, data, _ := nsq.UnpackResponse(resp)
   520  	t.Logf("frameType: %d, data: %s", frameType, data)
   521  	test.Equal(t, frameTypeResponse, frameType)
   522  	test.Equal(t, []byte("OK"), data)
   523  
   524  	time.Sleep(25 * time.Millisecond)
   525  
   526  	ch := nsqd.GetTopic(topicName).GetChannel("ch")
   527  	ch.deferredMutex.Lock()
   528  	numDef := len(ch.deferredMessages)
   529  	ch.deferredMutex.Unlock()
   530  	test.Equal(t, 1, numDef)
   531  	test.Equal(t, 1, int(atomic.LoadUint64(&ch.messageCount)))
   532  
   533  	// duration out of range
   534  	nsq.DeferredPublish(topicName, opts.MaxReqTimeout+100*time.Millisecond, make([]byte, 100)).WriteTo(conn)
   535  	resp, _ = nsq.ReadResponse(conn)
   536  	frameType, data, _ = nsq.UnpackResponse(resp)
   537  	t.Logf("frameType: %d, data: %s", frameType, data)
   538  	test.Equal(t, frameTypeError, frameType)
   539  	test.Equal(t, "E_INVALID DPUB timeout 3600100 out of range 0-3600000", string(data))
   540  }
   541  
   542  func TestUnixSocketTouch(t *testing.T) {
   543  	opts := NewOptions()
   544  	opts.Logger = test.NewTestLogger(t)
   545  	opts.LogLevel = LOG_DEBUG
   546  	opts.MsgTimeout = 150 * time.Millisecond
   547  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   548  	defer os.RemoveAll(opts.DataPath)
   549  	defer nsqd.Exit()
   550  
   551  	topicName := "test_touch" + strconv.Itoa(int(time.Now().Unix()))
   552  
   553  	conn, err := mustUnixSocketConnectNSQD(addr)
   554  	test.Nil(t, err)
   555  	defer conn.Close()
   556  
   557  	identify(t, conn, nil, frameTypeResponse)
   558  	sub(t, conn, topicName, "ch")
   559  
   560  	topic := nsqd.GetTopic(topicName)
   561  	channel := topic.GetChannel("ch")
   562  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
   563  	topic.PutMessage(msg)
   564  
   565  	_, err = nsq.Ready(1).WriteTo(conn)
   566  	test.Nil(t, err)
   567  
   568  	resp, err := nsq.ReadResponse(conn)
   569  	test.Nil(t, err)
   570  	frameType, data, _ := nsq.UnpackResponse(resp)
   571  	msgOut, _ := decodeMessage(data)
   572  	test.Equal(t, frameTypeMessage, frameType)
   573  	test.Equal(t, msg.ID, msgOut.ID)
   574  
   575  	time.Sleep(75 * time.Millisecond)
   576  
   577  	_, err = nsq.Touch(nsq.MessageID(msg.ID)).WriteTo(conn)
   578  	test.Nil(t, err)
   579  
   580  	time.Sleep(75 * time.Millisecond)
   581  
   582  	_, err = nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(conn)
   583  	test.Nil(t, err)
   584  
   585  	test.Equal(t, uint64(0), channel.timeoutCount)
   586  }
   587  
   588  func TestUnixSocketMaxRdyCount(t *testing.T) {
   589  	opts := NewOptions()
   590  	opts.Logger = test.NewTestLogger(t)
   591  	opts.LogLevel = LOG_DEBUG
   592  	opts.MaxRdyCount = 50
   593  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   594  	defer os.RemoveAll(opts.DataPath)
   595  	defer nsqd.Exit()
   596  
   597  	topicName := "test_max_rdy_count" + strconv.Itoa(int(time.Now().Unix()))
   598  
   599  	conn, err := mustUnixSocketConnectNSQD(addr)
   600  	test.Nil(t, err)
   601  	defer conn.Close()
   602  
   603  	topic := nsqd.GetTopic(topicName)
   604  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
   605  	topic.PutMessage(msg)
   606  
   607  	data := identify(t, conn, nil, frameTypeResponse)
   608  	r := struct {
   609  		MaxRdyCount int64 `json:"max_rdy_count"`
   610  	}{}
   611  	err = json.Unmarshal(data, &r)
   612  	test.Nil(t, err)
   613  	test.Equal(t, int64(50), r.MaxRdyCount)
   614  	sub(t, conn, topicName, "ch")
   615  
   616  	_, err = nsq.Ready(int(opts.MaxRdyCount)).WriteTo(conn)
   617  	test.Nil(t, err)
   618  
   619  	resp, err := nsq.ReadResponse(conn)
   620  	test.Nil(t, err)
   621  	frameType, data, _ := nsq.UnpackResponse(resp)
   622  	msgOut, _ := decodeMessage(data)
   623  	test.Equal(t, frameTypeMessage, frameType)
   624  	test.Equal(t, msg.ID, msgOut.ID)
   625  
   626  	_, err = nsq.Ready(int(opts.MaxRdyCount) + 1).WriteTo(conn)
   627  	test.Nil(t, err)
   628  
   629  	resp, err = nsq.ReadResponse(conn)
   630  	test.Nil(t, err)
   631  	frameType, data, _ = nsq.UnpackResponse(resp)
   632  	test.Equal(t, int32(1), frameType)
   633  	test.Equal(t, "E_INVALID RDY count 51 out of range 0-50", string(data))
   634  }
   635  
   636  func TestUnixSocketFatalError(t *testing.T) {
   637  	opts := NewOptions()
   638  	opts.Logger = test.NewTestLogger(t)
   639  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   640  	defer os.RemoveAll(opts.DataPath)
   641  	defer nsqd.Exit()
   642  
   643  	conn, err := mustUnixSocketConnectNSQD(addr)
   644  	test.Nil(t, err)
   645  	defer conn.Close()
   646  
   647  	_, err = conn.Write([]byte("ASDF\n"))
   648  	test.Nil(t, err)
   649  
   650  	resp, err := nsq.ReadResponse(conn)
   651  	test.Nil(t, err)
   652  	frameType, data, _ := nsq.UnpackResponse(resp)
   653  	test.Equal(t, int32(1), frameType)
   654  	test.Equal(t, "E_INVALID invalid command ASDF", string(data))
   655  
   656  	_, err = nsq.ReadResponse(conn)
   657  	test.NotNil(t, err)
   658  }
   659  
   660  func TestUnixSocketOutputBuffering(t *testing.T) {
   661  	opts := NewOptions()
   662  	opts.Logger = test.NewTestLogger(t)
   663  	opts.LogLevel = LOG_DEBUG
   664  	opts.MaxOutputBufferSize = 512 * 1024
   665  	opts.MaxOutputBufferTimeout = time.Second
   666  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   667  	defer os.RemoveAll(opts.DataPath)
   668  	defer nsqd.Exit()
   669  
   670  	topicName := "test_output_buffering" + strconv.Itoa(int(time.Now().Unix()))
   671  
   672  	conn, err := mustUnixSocketConnectNSQD(addr)
   673  	test.Nil(t, err)
   674  	defer conn.Close()
   675  
   676  	outputBufferSize := 256 * 1024
   677  	outputBufferTimeout := 500
   678  
   679  	topic := nsqd.GetTopic(topicName)
   680  	msg := NewMessage(topic.GenerateID(), make([]byte, outputBufferSize-1024))
   681  	topic.PutMessage(msg)
   682  
   683  	start := time.Now()
   684  	data := identify(t, conn, map[string]interface{}{
   685  		"output_buffer_size":    outputBufferSize,
   686  		"output_buffer_timeout": outputBufferTimeout,
   687  	}, frameTypeResponse)
   688  	var decoded map[string]interface{}
   689  	json.Unmarshal(data, &decoded)
   690  	v, ok := decoded["output_buffer_size"]
   691  	test.Equal(t, true, ok)
   692  	test.Equal(t, outputBufferSize, int(v.(float64)))
   693  	v = decoded["output_buffer_timeout"]
   694  	test.Equal(t, outputBufferTimeout, int(v.(float64)))
   695  	sub(t, conn, topicName, "ch")
   696  
   697  	_, err = nsq.Ready(10).WriteTo(conn)
   698  	test.Nil(t, err)
   699  
   700  	resp, err := nsq.ReadResponse(conn)
   701  	test.Nil(t, err)
   702  	end := time.Now()
   703  
   704  	test.Equal(t, true, int(end.Sub(start)/time.Millisecond) >= outputBufferTimeout)
   705  
   706  	frameType, data, _ := nsq.UnpackResponse(resp)
   707  	msgOut, _ := decodeMessage(data)
   708  	test.Equal(t, frameTypeMessage, frameType)
   709  	test.Equal(t, msg.ID, msgOut.ID)
   710  }
   711  
   712  func TestUnixSocketOutputBufferingValidity(t *testing.T) {
   713  	opts := NewOptions()
   714  	opts.Logger = test.NewTestLogger(t)
   715  	opts.LogLevel = LOG_DEBUG
   716  	opts.MaxOutputBufferSize = 512 * 1024
   717  	opts.MaxOutputBufferTimeout = time.Second
   718  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   719  	defer os.RemoveAll(opts.DataPath)
   720  	defer nsqd.Exit()
   721  
   722  	conn, err := mustUnixSocketConnectNSQD(addr)
   723  	test.Nil(t, err)
   724  	defer conn.Close()
   725  
   726  	identify(t, conn, map[string]interface{}{
   727  		"output_buffer_size":    512 * 1024,
   728  		"output_buffer_timeout": 1000,
   729  	}, frameTypeResponse)
   730  	identify(t, conn, map[string]interface{}{
   731  		"output_buffer_size":    -1,
   732  		"output_buffer_timeout": -1,
   733  	}, frameTypeResponse)
   734  	identify(t, conn, map[string]interface{}{
   735  		"output_buffer_size":    0,
   736  		"output_buffer_timeout": 0,
   737  	}, frameTypeResponse)
   738  	data := identify(t, conn, map[string]interface{}{
   739  		"output_buffer_size":    512*1024 + 1,
   740  		"output_buffer_timeout": 0,
   741  	}, frameTypeError)
   742  	test.Equal(t, fmt.Sprintf("E_BAD_BODY IDENTIFY output buffer size (%d) is invalid", 512*1024+1), string(data))
   743  
   744  	conn, err = mustUnixSocketConnectNSQD(addr)
   745  	test.Nil(t, err)
   746  	defer conn.Close()
   747  
   748  	data = identify(t, conn, map[string]interface{}{
   749  		"output_buffer_size":    0,
   750  		"output_buffer_timeout": 1001,
   751  	}, frameTypeError)
   752  	test.Equal(t, "E_BAD_BODY IDENTIFY output buffer timeout (1001) is invalid", string(data))
   753  }
   754  
   755  func TestUnixSocketTLS(t *testing.T) {
   756  	opts := NewOptions()
   757  	opts.Logger = test.NewTestLogger(t)
   758  	opts.LogLevel = LOG_DEBUG
   759  	opts.TLSCert = "./test/certs/server.pem"
   760  	opts.TLSKey = "./test/certs/server.key"
   761  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   762  	defer os.RemoveAll(opts.DataPath)
   763  	defer nsqd.Exit()
   764  
   765  	conn, err := mustUnixSocketConnectNSQD(addr)
   766  	test.Nil(t, err)
   767  	defer conn.Close()
   768  
   769  	data := identify(t, conn, map[string]interface{}{
   770  		"tls_v1": true,
   771  	}, frameTypeResponse)
   772  	r := struct {
   773  		TLSv1 bool `json:"tls_v1"`
   774  	}{}
   775  	err = json.Unmarshal(data, &r)
   776  	test.Nil(t, err)
   777  	test.Equal(t, true, r.TLSv1)
   778  
   779  	tlsConfig := &tls.Config{
   780  		InsecureSkipVerify: true,
   781  	}
   782  	tlsConn := tls.Client(conn, tlsConfig)
   783  
   784  	err = tlsConn.Handshake()
   785  	test.Nil(t, err)
   786  
   787  	resp, _ := nsq.ReadResponse(tlsConn)
   788  	frameType, data, _ := nsq.UnpackResponse(resp)
   789  	t.Logf("frameType: %d, data: %s", frameType, data)
   790  	test.Equal(t, frameTypeResponse, frameType)
   791  	test.Equal(t, []byte("OK"), data)
   792  }
   793  
   794  func TestUnixSocketTLSRequired(t *testing.T) {
   795  	opts := NewOptions()
   796  	opts.Logger = test.NewTestLogger(t)
   797  	opts.LogLevel = LOG_DEBUG
   798  	opts.TLSCert = "./test/certs/server.pem"
   799  	opts.TLSKey = "./test/certs/server.key"
   800  	opts.TLSRequired = TLSRequiredExceptHTTP
   801  
   802  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   803  	defer os.RemoveAll(opts.DataPath)
   804  	defer nsqd.Exit()
   805  
   806  	topicName := "test_tls_required" + strconv.Itoa(int(time.Now().Unix()))
   807  
   808  	conn, err := mustUnixSocketConnectNSQD(addr)
   809  	test.Nil(t, err)
   810  	defer conn.Close()
   811  
   812  	subFail(t, conn, topicName, "ch")
   813  
   814  	conn, err = mustUnixSocketConnectNSQD(addr)
   815  	test.Nil(t, err)
   816  	defer conn.Close()
   817  
   818  	data := identify(t, conn, map[string]interface{}{
   819  		"tls_v1": true,
   820  	}, frameTypeResponse)
   821  	r := struct {
   822  		TLSv1 bool `json:"tls_v1"`
   823  	}{}
   824  	err = json.Unmarshal(data, &r)
   825  	test.Nil(t, err)
   826  	test.Equal(t, true, r.TLSv1)
   827  
   828  	tlsConfig := &tls.Config{
   829  		InsecureSkipVerify: true,
   830  	}
   831  	tlsConn := tls.Client(conn, tlsConfig)
   832  
   833  	err = tlsConn.Handshake()
   834  	test.Nil(t, err)
   835  
   836  	resp, _ := nsq.ReadResponse(tlsConn)
   837  	frameType, data, _ := nsq.UnpackResponse(resp)
   838  	t.Logf("frameType: %d, data: %s", frameType, data)
   839  	test.Equal(t, frameTypeResponse, frameType)
   840  	test.Equal(t, []byte("OK"), data)
   841  }
   842  
   843  func TestUnixSocketTLSAuthRequire(t *testing.T) {
   844  	opts := NewOptions()
   845  	opts.Logger = test.NewTestLogger(t)
   846  	opts.LogLevel = LOG_DEBUG
   847  	opts.TLSCert = "./test/certs/server.pem"
   848  	opts.TLSKey = "./test/certs/server.key"
   849  	opts.TLSClientAuthPolicy = "require"
   850  
   851  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   852  	defer os.RemoveAll(opts.DataPath)
   853  	defer nsqd.Exit()
   854  
   855  	// No Certs
   856  	conn, err := mustUnixSocketConnectNSQD(addr)
   857  	test.Nil(t, err)
   858  	defer conn.Close()
   859  
   860  	data := identify(t, conn, map[string]interface{}{
   861  		"tls_v1": true,
   862  	}, frameTypeResponse)
   863  	r := struct {
   864  		TLSv1 bool `json:"tls_v1"`
   865  	}{}
   866  	err = json.Unmarshal(data, &r)
   867  	test.Nil(t, err)
   868  	test.Equal(t, true, r.TLSv1)
   869  	tlsConfig := &tls.Config{
   870  		InsecureSkipVerify: true,
   871  	}
   872  	tlsConn := tls.Client(conn, tlsConfig)
   873  	_, err = nsq.ReadResponse(tlsConn)
   874  	test.NotNil(t, err)
   875  
   876  	// With Unsigned Cert
   877  	conn, err = mustUnixSocketConnectNSQD(addr)
   878  	test.Nil(t, err)
   879  	defer conn.Close()
   880  
   881  	data = identify(t, conn, map[string]interface{}{
   882  		"tls_v1": true,
   883  	}, frameTypeResponse)
   884  	r = struct {
   885  		TLSv1 bool `json:"tls_v1"`
   886  	}{}
   887  	err = json.Unmarshal(data, &r)
   888  	test.Nil(t, err)
   889  	test.Equal(t, true, r.TLSv1)
   890  
   891  	cert, err := tls.LoadX509KeyPair("./test/certs/cert.pem", "./test/certs/key.pem")
   892  	test.Nil(t, err)
   893  	tlsConfig = &tls.Config{
   894  		Certificates:       []tls.Certificate{cert},
   895  		InsecureSkipVerify: true,
   896  	}
   897  	tlsConn = tls.Client(conn, tlsConfig)
   898  	err = tlsConn.Handshake()
   899  	test.Nil(t, err)
   900  
   901  	resp, _ := nsq.ReadResponse(tlsConn)
   902  	frameType, data, _ := nsq.UnpackResponse(resp)
   903  	t.Logf("frameType: %d, data: %s", frameType, data)
   904  	test.Equal(t, frameTypeResponse, frameType)
   905  	test.Equal(t, []byte("OK"), data)
   906  
   907  }
   908  
   909  func TestUnixSocketTLSAuthRequireVerify(t *testing.T) {
   910  	opts := NewOptions()
   911  	opts.Logger = test.NewTestLogger(t)
   912  	opts.LogLevel = LOG_DEBUG
   913  	opts.TLSCert = "./test/certs/server.pem"
   914  	opts.TLSKey = "./test/certs/server.key"
   915  	opts.TLSRootCAFile = "./test/certs/ca.pem"
   916  	opts.TLSClientAuthPolicy = "require-verify"
   917  
   918  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
   919  	defer os.RemoveAll(opts.DataPath)
   920  	defer nsqd.Exit()
   921  
   922  	// with no cert
   923  	conn, err := mustUnixSocketConnectNSQD(addr)
   924  	test.Nil(t, err)
   925  	defer conn.Close()
   926  
   927  	data := identify(t, conn, map[string]interface{}{
   928  		"tls_v1": true,
   929  	}, frameTypeResponse)
   930  	r := struct {
   931  		TLSv1 bool `json:"tls_v1"`
   932  	}{}
   933  	err = json.Unmarshal(data, &r)
   934  	test.Nil(t, err)
   935  	test.Equal(t, true, r.TLSv1)
   936  	tlsConfig := &tls.Config{
   937  		InsecureSkipVerify: true,
   938  	}
   939  	tlsConn := tls.Client(conn, tlsConfig)
   940  	_, err = nsq.ReadResponse(tlsConn)
   941  	test.NotNil(t, err)
   942  
   943  	// with invalid cert
   944  	conn, err = mustUnixSocketConnectNSQD(addr)
   945  	test.Nil(t, err)
   946  	defer conn.Close()
   947  
   948  	data = identify(t, conn, map[string]interface{}{
   949  		"tls_v1": true,
   950  	}, frameTypeResponse)
   951  	r = struct {
   952  		TLSv1 bool `json:"tls_v1"`
   953  	}{}
   954  	err = json.Unmarshal(data, &r)
   955  	test.Nil(t, err)
   956  	test.Equal(t, true, r.TLSv1)
   957  	cert, err := tls.LoadX509KeyPair("./test/certs/cert.pem", "./test/certs/key.pem")
   958  	test.Nil(t, err)
   959  	tlsConfig = &tls.Config{
   960  		Certificates:       []tls.Certificate{cert},
   961  		InsecureSkipVerify: true,
   962  	}
   963  	tlsConn = tls.Client(conn, tlsConfig)
   964  	_, err = nsq.ReadResponse(tlsConn)
   965  	test.NotNil(t, err)
   966  
   967  	// with valid cert
   968  	conn, err = mustUnixSocketConnectNSQD(addr)
   969  	test.Nil(t, err)
   970  	defer conn.Close()
   971  
   972  	data = identify(t, conn, map[string]interface{}{
   973  		"tls_v1": true,
   974  	}, frameTypeResponse)
   975  	r = struct {
   976  		TLSv1 bool `json:"tls_v1"`
   977  	}{}
   978  	err = json.Unmarshal(data, &r)
   979  	test.Nil(t, err)
   980  	test.Equal(t, true, r.TLSv1)
   981  	cert, err = tls.LoadX509KeyPair("./test/certs/client.pem", "./test/certs/client.key")
   982  	test.Nil(t, err)
   983  	tlsConfig = &tls.Config{
   984  		Certificates:       []tls.Certificate{cert},
   985  		InsecureSkipVerify: true,
   986  	}
   987  	tlsConn = tls.Client(conn, tlsConfig)
   988  	err = tlsConn.Handshake()
   989  	test.Nil(t, err)
   990  
   991  	resp, _ := nsq.ReadResponse(tlsConn)
   992  	frameType, data, _ := nsq.UnpackResponse(resp)
   993  	t.Logf("frameType: %d, data: %s", frameType, data)
   994  	test.Equal(t, frameTypeResponse, frameType)
   995  	test.Equal(t, []byte("OK"), data)
   996  }
   997  
   998  func TestUnixSocketDeflate(t *testing.T) {
   999  	opts := NewOptions()
  1000  	opts.Logger = test.NewTestLogger(t)
  1001  	opts.LogLevel = LOG_DEBUG
  1002  	opts.DeflateEnabled = true
  1003  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1004  	defer os.RemoveAll(opts.DataPath)
  1005  	defer nsqd.Exit()
  1006  
  1007  	conn, err := mustUnixSocketConnectNSQD(addr)
  1008  	test.Nil(t, err)
  1009  	defer conn.Close()
  1010  
  1011  	data := identify(t, conn, map[string]interface{}{
  1012  		"deflate": true,
  1013  	}, frameTypeResponse)
  1014  	r := struct {
  1015  		Deflate bool `json:"deflate"`
  1016  	}{}
  1017  	err = json.Unmarshal(data, &r)
  1018  	test.Nil(t, err)
  1019  	test.Equal(t, true, r.Deflate)
  1020  
  1021  	compressConn := flate.NewReader(conn)
  1022  	resp, _ := nsq.ReadResponse(compressConn)
  1023  	frameType, data, _ := nsq.UnpackResponse(resp)
  1024  	t.Logf("frameType: %d, data: %s", frameType, data)
  1025  	test.Equal(t, frameTypeResponse, frameType)
  1026  	test.Equal(t, []byte("OK"), data)
  1027  }
  1028  
  1029  func TestUnixSocketSnappy(t *testing.T) {
  1030  	opts := NewOptions()
  1031  	opts.Logger = test.NewTestLogger(t)
  1032  	opts.LogLevel = LOG_DEBUG
  1033  	opts.SnappyEnabled = true
  1034  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1035  	defer os.RemoveAll(opts.DataPath)
  1036  	defer nsqd.Exit()
  1037  
  1038  	conn, err := mustUnixSocketConnectNSQD(addr)
  1039  	test.Nil(t, err)
  1040  	defer conn.Close()
  1041  
  1042  	data := identify(t, conn, map[string]interface{}{
  1043  		"snappy": true,
  1044  	}, frameTypeResponse)
  1045  	r := struct {
  1046  		Snappy bool `json:"snappy"`
  1047  	}{}
  1048  	err = json.Unmarshal(data, &r)
  1049  	test.Nil(t, err)
  1050  	test.Equal(t, true, r.Snappy)
  1051  
  1052  	compressConn := snappy.NewReader(conn)
  1053  	resp, _ := nsq.ReadResponse(compressConn)
  1054  	frameType, data, _ := nsq.UnpackResponse(resp)
  1055  	t.Logf("frameType: %d, data: %s", frameType, data)
  1056  	test.Equal(t, frameTypeResponse, frameType)
  1057  	test.Equal(t, []byte("OK"), data)
  1058  
  1059  	msgBody := make([]byte, 128000)
  1060  	//lint:ignore SA1019 NewWriter is deprecated by NewBufferedWriter, but we don't want to buffer
  1061  	w := snappy.NewWriter(conn)
  1062  
  1063  	rw := readWriter{compressConn, w}
  1064  
  1065  	topicName := "test_snappy" + strconv.Itoa(int(time.Now().Unix()))
  1066  	sub(t, rw, topicName, "ch")
  1067  
  1068  	_, err = nsq.Ready(1).WriteTo(rw)
  1069  	test.Nil(t, err)
  1070  
  1071  	topic := nsqd.GetTopic(topicName)
  1072  	msg := NewMessage(topic.GenerateID(), msgBody)
  1073  	topic.PutMessage(msg)
  1074  
  1075  	resp, _ = nsq.ReadResponse(compressConn)
  1076  	frameType, data, _ = nsq.UnpackResponse(resp)
  1077  	msgOut, _ := decodeMessage(data)
  1078  	test.Equal(t, frameTypeMessage, frameType)
  1079  	test.Equal(t, msg.ID, msgOut.ID)
  1080  	test.Equal(t, msg.Body, msgOut.Body)
  1081  }
  1082  
  1083  func TestUnixSocketTLSDeflate(t *testing.T) {
  1084  	opts := NewOptions()
  1085  	opts.Logger = test.NewTestLogger(t)
  1086  	opts.LogLevel = LOG_DEBUG
  1087  	opts.DeflateEnabled = true
  1088  	opts.TLSCert = "./test/certs/cert.pem"
  1089  	opts.TLSKey = "./test/certs/key.pem"
  1090  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1091  	defer os.RemoveAll(opts.DataPath)
  1092  	defer nsqd.Exit()
  1093  
  1094  	conn, err := mustUnixSocketConnectNSQD(addr)
  1095  	test.Nil(t, err)
  1096  	defer conn.Close()
  1097  
  1098  	data := identify(t, conn, map[string]interface{}{
  1099  		"tls_v1":  true,
  1100  		"deflate": true,
  1101  	}, frameTypeResponse)
  1102  	r := struct {
  1103  		TLSv1   bool `json:"tls_v1"`
  1104  		Deflate bool `json:"deflate"`
  1105  	}{}
  1106  	err = json.Unmarshal(data, &r)
  1107  	test.Nil(t, err)
  1108  	test.Equal(t, true, r.TLSv1)
  1109  	test.Equal(t, true, r.Deflate)
  1110  
  1111  	tlsConfig := &tls.Config{
  1112  		InsecureSkipVerify: true,
  1113  	}
  1114  	tlsConn := tls.Client(conn, tlsConfig)
  1115  
  1116  	err = tlsConn.Handshake()
  1117  	test.Nil(t, err)
  1118  
  1119  	resp, _ := nsq.ReadResponse(tlsConn)
  1120  	frameType, data, _ := nsq.UnpackResponse(resp)
  1121  	t.Logf("frameType: %d, data: %s", frameType, data)
  1122  	test.Equal(t, frameTypeResponse, frameType)
  1123  	test.Equal(t, []byte("OK"), data)
  1124  
  1125  	compressConn := flate.NewReader(tlsConn)
  1126  
  1127  	resp, _ = nsq.ReadResponse(compressConn)
  1128  	frameType, data, _ = nsq.UnpackResponse(resp)
  1129  	t.Logf("frameType: %d, data: %s", frameType, data)
  1130  	test.Equal(t, frameTypeResponse, frameType)
  1131  	test.Equal(t, []byte("OK"), data)
  1132  }
  1133  
  1134  func TestUnixSocketSampling(t *testing.T) {
  1135  	rand.Seed(time.Now().UTC().UnixNano())
  1136  
  1137  	num := 10000
  1138  	sampleRate := 42
  1139  	slack := 5
  1140  
  1141  	opts := NewOptions()
  1142  	opts.Logger = test.NewTestLogger(t)
  1143  	opts.LogLevel = LOG_DEBUG
  1144  	opts.MaxRdyCount = int64(num)
  1145  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1146  	defer os.RemoveAll(opts.DataPath)
  1147  	defer nsqd.Exit()
  1148  
  1149  	conn, err := mustUnixSocketConnectNSQD(addr)
  1150  	test.Nil(t, err)
  1151  	defer conn.Close()
  1152  
  1153  	data := identify(t, conn, map[string]interface{}{
  1154  		"sample_rate": int32(sampleRate),
  1155  	}, frameTypeResponse)
  1156  	r := struct {
  1157  		SampleRate int32 `json:"sample_rate"`
  1158  	}{}
  1159  	err = json.Unmarshal(data, &r)
  1160  	test.Nil(t, err)
  1161  	test.Equal(t, int32(sampleRate), r.SampleRate)
  1162  
  1163  	topicName := "test_sampling" + strconv.Itoa(int(time.Now().Unix()))
  1164  	topic := nsqd.GetTopic(topicName)
  1165  	for i := 0; i < num; i++ {
  1166  		msg := NewMessage(topic.GenerateID(), []byte("test body"))
  1167  		topic.PutMessage(msg)
  1168  	}
  1169  	channel := topic.GetChannel("ch")
  1170  
  1171  	// let the topic drain into the channel
  1172  	time.Sleep(50 * time.Millisecond)
  1173  
  1174  	sub(t, conn, topicName, "ch")
  1175  	_, err = nsq.Ready(num).WriteTo(conn)
  1176  	test.Nil(t, err)
  1177  
  1178  	go func() {
  1179  		for {
  1180  			_, err := nsq.ReadResponse(conn)
  1181  			if err != nil {
  1182  				return
  1183  			}
  1184  		}
  1185  	}()
  1186  
  1187  	doneChan := make(chan int)
  1188  	go func() {
  1189  		for {
  1190  			if channel.Depth() == 0 {
  1191  				close(doneChan)
  1192  				return
  1193  			}
  1194  			time.Sleep(5 * time.Millisecond)
  1195  		}
  1196  	}()
  1197  	<-doneChan
  1198  
  1199  	channel.inFlightMutex.Lock()
  1200  	numInFlight := len(channel.inFlightMessages)
  1201  	channel.inFlightMutex.Unlock()
  1202  
  1203  	test.Equal(t, true, numInFlight <= int(float64(num)*float64(sampleRate+slack)/100.0))
  1204  	test.Equal(t, true, numInFlight >= int(float64(num)*float64(sampleRate-slack)/100.0))
  1205  }
  1206  
  1207  func TestUnixSocketTLSSnappy(t *testing.T) {
  1208  	opts := NewOptions()
  1209  	opts.Logger = test.NewTestLogger(t)
  1210  	opts.LogLevel = LOG_DEBUG
  1211  	opts.SnappyEnabled = true
  1212  	opts.TLSCert = "./test/certs/cert.pem"
  1213  	opts.TLSKey = "./test/certs/key.pem"
  1214  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1215  	defer os.RemoveAll(opts.DataPath)
  1216  	defer nsqd.Exit()
  1217  
  1218  	conn, err := mustUnixSocketConnectNSQD(addr)
  1219  	test.Nil(t, err)
  1220  	defer conn.Close()
  1221  
  1222  	data := identify(t, conn, map[string]interface{}{
  1223  		"tls_v1": true,
  1224  		"snappy": true,
  1225  	}, frameTypeResponse)
  1226  	r := struct {
  1227  		TLSv1  bool `json:"tls_v1"`
  1228  		Snappy bool `json:"snappy"`
  1229  	}{}
  1230  	err = json.Unmarshal(data, &r)
  1231  	test.Nil(t, err)
  1232  	test.Equal(t, true, r.TLSv1)
  1233  	test.Equal(t, true, r.Snappy)
  1234  
  1235  	tlsConfig := &tls.Config{
  1236  		InsecureSkipVerify: true,
  1237  	}
  1238  	tlsConn := tls.Client(conn, tlsConfig)
  1239  
  1240  	err = tlsConn.Handshake()
  1241  	test.Nil(t, err)
  1242  
  1243  	resp, _ := nsq.ReadResponse(tlsConn)
  1244  	frameType, data, _ := nsq.UnpackResponse(resp)
  1245  	t.Logf("frameType: %d, data: %s", frameType, data)
  1246  	test.Equal(t, frameTypeResponse, frameType)
  1247  	test.Equal(t, []byte("OK"), data)
  1248  
  1249  	compressConn := snappy.NewReader(tlsConn)
  1250  
  1251  	resp, _ = nsq.ReadResponse(compressConn)
  1252  	frameType, data, _ = nsq.UnpackResponse(resp)
  1253  	t.Logf("frameType: %d, data: %s", frameType, data)
  1254  	test.Equal(t, frameTypeResponse, frameType)
  1255  	test.Equal(t, []byte("OK"), data)
  1256  }
  1257  
  1258  func TestUnixSocketClientMsgTimeout(t *testing.T) {
  1259  	opts := NewOptions()
  1260  	opts.Logger = test.NewTestLogger(t)
  1261  	opts.LogLevel = LOG_DEBUG
  1262  	opts.QueueScanRefreshInterval = 100 * time.Millisecond
  1263  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1264  	defer os.RemoveAll(opts.DataPath)
  1265  	defer nsqd.Exit()
  1266  
  1267  	topicName := "test_cmsg_timeout" + strconv.Itoa(int(time.Now().Unix()))
  1268  	topic := nsqd.GetTopic(topicName)
  1269  	ch := topic.GetChannel("ch")
  1270  	msg := NewMessage(topic.GenerateID(), make([]byte, 100))
  1271  	topic.PutMessage(msg)
  1272  
  1273  	// without this the race detector thinks there's a write
  1274  	// to msg.Attempts that races with the read in the protocol's messagePump...
  1275  	// it does not reflect a realistically possible condition
  1276  	topic.PutMessage(NewMessage(topic.GenerateID(), make([]byte, 100)))
  1277  
  1278  	conn, err := mustUnixSocketConnectNSQD(addr)
  1279  	test.Nil(t, err)
  1280  	defer conn.Close()
  1281  
  1282  	identify(t, conn, map[string]interface{}{
  1283  		"msg_timeout": 1000,
  1284  	}, frameTypeResponse)
  1285  	sub(t, conn, topicName, "ch")
  1286  
  1287  	test.Equal(t, 0, int(atomic.LoadUint64(&ch.timeoutCount)))
  1288  	test.Equal(t, 0, int(atomic.LoadUint64(&ch.requeueCount)))
  1289  
  1290  	_, err = nsq.Ready(1).WriteTo(conn)
  1291  	test.Nil(t, err)
  1292  
  1293  	resp, _ := nsq.ReadResponse(conn)
  1294  	_, data, _ := nsq.UnpackResponse(resp)
  1295  	msgOut, _ := decodeMessage(data)
  1296  	test.Equal(t, msg.ID, msgOut.ID)
  1297  	test.Equal(t, msg.Body, msgOut.Body)
  1298  
  1299  	_, err = nsq.Ready(0).WriteTo(conn)
  1300  	test.Nil(t, err)
  1301  
  1302  	time.Sleep(1150 * time.Millisecond)
  1303  
  1304  	test.Equal(t, 1, int(atomic.LoadUint64(&ch.timeoutCount)))
  1305  	test.Equal(t, 0, int(atomic.LoadUint64(&ch.requeueCount)))
  1306  
  1307  	_, err = nsq.Finish(nsq.MessageID(msgOut.ID)).WriteTo(conn)
  1308  	test.Nil(t, err)
  1309  
  1310  	resp, _ = nsq.ReadResponse(conn)
  1311  	frameType, data, _ := nsq.UnpackResponse(resp)
  1312  	test.Equal(t, frameTypeError, frameType)
  1313  	test.Equal(t, fmt.Sprintf("E_FIN_FAILED FIN %s failed ID not in flight", msgOut.ID),
  1314  		string(data))
  1315  }
  1316  
  1317  func TestUnixSocketBadFin(t *testing.T) {
  1318  	opts := NewOptions()
  1319  	opts.Logger = test.NewTestLogger(t)
  1320  	opts.LogLevel = LOG_DEBUG
  1321  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1322  	defer os.RemoveAll(opts.DataPath)
  1323  	defer nsqd.Exit()
  1324  
  1325  	conn, err := mustUnixSocketConnectNSQD(addr)
  1326  	test.Nil(t, err)
  1327  	defer conn.Close()
  1328  
  1329  	identify(t, conn, map[string]interface{}{}, frameTypeResponse)
  1330  	sub(t, conn, "test_fin", "ch")
  1331  
  1332  	fin := nsq.Finish(nsq.MessageID{})
  1333  	fin.Params[0] = []byte("")
  1334  	_, err = fin.WriteTo(conn)
  1335  	test.Nil(t, err)
  1336  
  1337  	resp, _ := nsq.ReadResponse(conn)
  1338  	frameType, data, _ := nsq.UnpackResponse(resp)
  1339  	test.Equal(t, frameTypeError, frameType)
  1340  	test.Equal(t, "E_INVALID invalid message ID", string(data))
  1341  }
  1342  
  1343  func TestUnixSocketReqTimeoutRange(t *testing.T) {
  1344  	opts := NewOptions()
  1345  	opts.Logger = test.NewTestLogger(t)
  1346  	opts.LogLevel = LOG_DEBUG
  1347  	opts.MaxReqTimeout = 1 * time.Minute
  1348  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1349  	defer os.RemoveAll(opts.DataPath)
  1350  	defer nsqd.Exit()
  1351  
  1352  	topicName := "test_req" + strconv.Itoa(int(time.Now().Unix()))
  1353  
  1354  	conn, err := mustUnixSocketConnectNSQD(addr)
  1355  	test.Nil(t, err)
  1356  	defer conn.Close()
  1357  
  1358  	identify(t, conn, nil, frameTypeResponse)
  1359  	sub(t, conn, topicName, "ch")
  1360  
  1361  	topic := nsqd.GetTopic(topicName)
  1362  	channel := topic.GetChannel("ch")
  1363  	msg := NewMessage(topic.GenerateID(), []byte("test body"))
  1364  	topic.PutMessage(msg)
  1365  
  1366  	_, err = nsq.Ready(1).WriteTo(conn)
  1367  	test.Nil(t, err)
  1368  
  1369  	resp, err := nsq.ReadResponse(conn)
  1370  	test.Nil(t, err)
  1371  	frameType, data, _ := nsq.UnpackResponse(resp)
  1372  	msgOut, _ := decodeMessage(data)
  1373  	test.Equal(t, frameTypeMessage, frameType)
  1374  	test.Equal(t, msg.ID, msgOut.ID)
  1375  
  1376  	_, err = nsq.Requeue(nsq.MessageID(msg.ID), -1).WriteTo(conn)
  1377  	test.Nil(t, err)
  1378  
  1379  	// It should be immediately available for another attempt
  1380  	resp, err = nsq.ReadResponse(conn)
  1381  	test.Nil(t, err)
  1382  	frameType, data, _ = nsq.UnpackResponse(resp)
  1383  	msgOut, _ = decodeMessage(data)
  1384  	test.Equal(t, frameTypeMessage, frameType)
  1385  	test.Equal(t, msg.ID, msgOut.ID)
  1386  
  1387  	// The priority (processing time) should be >= this
  1388  	minTs := time.Now().Add(opts.MaxReqTimeout).UnixNano()
  1389  
  1390  	_, err = nsq.Requeue(nsq.MessageID(msg.ID), opts.MaxReqTimeout*2).WriteTo(conn)
  1391  	test.Nil(t, err)
  1392  
  1393  	time.Sleep(100 * time.Millisecond)
  1394  
  1395  	channel.deferredMutex.Lock()
  1396  	pqItem := channel.deferredMessages[msg.ID]
  1397  	channel.deferredMutex.Unlock()
  1398  
  1399  	test.NotNil(t, pqItem)
  1400  	test.Equal(t, true, pqItem.Priority >= minTs)
  1401  }
  1402  
  1403  func TestUnixSocketIOLoopReturnsClientErrWhenSendFails(t *testing.T) {
  1404  	fakeConn := test.NewFakeNetConn()
  1405  	fakeConn.WriteFunc = func(b []byte) (int, error) {
  1406  		return 0, errors.New("write error")
  1407  	}
  1408  
  1409  	testUnixSocketIOLoopReturnsClientErr(t, fakeConn)
  1410  }
  1411  
  1412  func TestUnixSocketIOLoopReturnsClientErrWhenSendSucceeds(t *testing.T) {
  1413  	fakeConn := test.NewFakeNetConn()
  1414  	fakeConn.WriteFunc = func(b []byte) (int, error) {
  1415  		return len(b), nil
  1416  	}
  1417  
  1418  	testUnixSocketIOLoopReturnsClientErr(t, fakeConn)
  1419  }
  1420  
  1421  func testUnixSocketIOLoopReturnsClientErr(t *testing.T, fakeConn test.FakeNetConn) {
  1422  	fakeConn.ReadFunc = func(b []byte) (int, error) {
  1423  		return copy(b, []byte("INVALID_COMMAND\n")), nil
  1424  	}
  1425  
  1426  	opts := NewOptions()
  1427  	opts.Logger = test.NewTestLogger(t)
  1428  	opts.LogLevel = LOG_DEBUG
  1429  
  1430  	nsqd, err := New(opts)
  1431  	test.Nil(t, err)
  1432  	prot := &protocolV2{nsqd: nsqd}
  1433  	defer prot.nsqd.Exit()
  1434  
  1435  	client := prot.NewClient(fakeConn)
  1436  	err = prot.IOLoop(client)
  1437  	test.NotNil(t, err)
  1438  	test.Equal(t, "E_INVALID invalid command INVALID_COMMAND", err.Error())
  1439  	test.NotNil(t, err.(*protocol.FatalClientErr))
  1440  }
  1441  
  1442  func BenchmarkUnixSocketProtocolV2Exec(b *testing.B) {
  1443  	b.StopTimer()
  1444  	opts := NewOptions()
  1445  	opts.Logger = test.NewTestLogger(b)
  1446  	nsqd, _ := New(opts)
  1447  	p := &protocolV2{nsqd}
  1448  	c := newClientV2(0, nil, nsqd)
  1449  	params := [][]byte{[]byte("NOP")}
  1450  	b.StartTimer()
  1451  
  1452  	for i := 0; i < b.N; i++ {
  1453  		p.Exec(c, params)
  1454  	}
  1455  }
  1456  
  1457  func benchmarkUnixSocketProtocolV2PubMultiTopic(b *testing.B, numTopics int) {
  1458  	var wg sync.WaitGroup
  1459  	b.StopTimer()
  1460  	opts := NewOptions()
  1461  	size := 200
  1462  	batchSize := int(opts.MaxBodySize) / (size + 4)
  1463  	opts.Logger = test.NewTestLogger(b)
  1464  	opts.MemQueueSize = int64(b.N)
  1465  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1466  	defer os.RemoveAll(opts.DataPath)
  1467  	msg := make([]byte, size)
  1468  	batch := make([][]byte, batchSize)
  1469  	for i := range batch {
  1470  		batch[i] = msg
  1471  	}
  1472  	b.SetBytes(int64(len(msg)))
  1473  	b.StartTimer()
  1474  
  1475  	for j := 0; j < numTopics; j++ {
  1476  		topicName := fmt.Sprintf("bench_v2_pub_multi_topic_%d_%d", j, time.Now().Unix())
  1477  		wg.Add(1)
  1478  		go func() {
  1479  			conn, err := mustUnixSocketConnectNSQD(addr)
  1480  			if err != nil {
  1481  				panic(err.Error())
  1482  			}
  1483  			rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
  1484  
  1485  			num := b.N / numTopics / batchSize
  1486  			wg.Add(1)
  1487  			go func() {
  1488  				for i := 0; i < num; i++ {
  1489  					cmd, _ := nsq.MultiPublish(topicName, batch)
  1490  					_, err := cmd.WriteTo(rw)
  1491  					if err != nil {
  1492  						panic(err.Error())
  1493  					}
  1494  					err = rw.Flush()
  1495  					if err != nil {
  1496  						panic(err.Error())
  1497  					}
  1498  				}
  1499  				wg.Done()
  1500  			}()
  1501  			wg.Add(1)
  1502  			go func() {
  1503  				for i := 0; i < num; i++ {
  1504  					resp, err := nsq.ReadResponse(rw)
  1505  					if err != nil {
  1506  						panic(err.Error())
  1507  					}
  1508  					_, data, _ := nsq.UnpackResponse(resp)
  1509  					if !bytes.Equal(data, []byte("OK")) {
  1510  						panic("invalid response")
  1511  					}
  1512  				}
  1513  				wg.Done()
  1514  			}()
  1515  			wg.Done()
  1516  		}()
  1517  	}
  1518  
  1519  	wg.Wait()
  1520  
  1521  	b.StopTimer()
  1522  	nsqd.Exit()
  1523  }
  1524  
  1525  func BenchmarkUnixSocketProtocolV2PubMultiTopic1(b *testing.B) {
  1526  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 1)
  1527  }
  1528  func BenchmarkUnixSocketkProtocolV2PubMultiTopic2(b *testing.B) {
  1529  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 2)
  1530  }
  1531  func BenchmarkUnixSocketProtocolV2PubMultiTopic4(b *testing.B) {
  1532  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 4)
  1533  }
  1534  func BenchmarkUnixSocketProtocolV2PubMultiTopic8(b *testing.B) {
  1535  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 8)
  1536  }
  1537  func BenchmarkUnixSocketProtocolV2PubMultiTopic16(b *testing.B) {
  1538  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 16)
  1539  }
  1540  func BenchmarkUnixSocketProtocolV2PubMultiTopic32(b *testing.B) {
  1541  	benchmarkUnixSocketProtocolV2PubMultiTopic(b, 32)
  1542  }
  1543  
  1544  func benchmarkUnixSocketProtocolV2Pub(b *testing.B, size int) {
  1545  	var wg sync.WaitGroup
  1546  	b.StopTimer()
  1547  	opts := NewOptions()
  1548  	batchSize := int(opts.MaxBodySize) / (size + 4)
  1549  	opts.Logger = test.NewTestLogger(b)
  1550  	opts.MemQueueSize = int64(b.N)
  1551  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1552  	defer os.RemoveAll(opts.DataPath)
  1553  	msg := make([]byte, size)
  1554  	batch := make([][]byte, batchSize)
  1555  	for i := range batch {
  1556  		batch[i] = msg
  1557  	}
  1558  	topicName := "bench_v2_pub" + strconv.Itoa(int(time.Now().Unix()))
  1559  	b.SetBytes(int64(len(msg)))
  1560  	b.StartTimer()
  1561  
  1562  	for j := 0; j < runtime.GOMAXPROCS(0); j++ {
  1563  		wg.Add(1)
  1564  		go func() {
  1565  			conn, err := mustUnixSocketConnectNSQD(addr)
  1566  			if err != nil {
  1567  				panic(err.Error())
  1568  			}
  1569  			rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
  1570  
  1571  			num := b.N / runtime.GOMAXPROCS(0) / batchSize
  1572  			wg.Add(1)
  1573  			go func() {
  1574  				for i := 0; i < num; i++ {
  1575  					cmd, _ := nsq.MultiPublish(topicName, batch)
  1576  					_, err := cmd.WriteTo(rw)
  1577  					if err != nil {
  1578  						panic(err.Error())
  1579  					}
  1580  					err = rw.Flush()
  1581  					if err != nil {
  1582  						panic(err.Error())
  1583  					}
  1584  				}
  1585  				wg.Done()
  1586  			}()
  1587  			wg.Add(1)
  1588  			go func() {
  1589  				for i := 0; i < num; i++ {
  1590  					resp, err := nsq.ReadResponse(rw)
  1591  					if err != nil {
  1592  						panic(err.Error())
  1593  					}
  1594  					_, data, _ := nsq.UnpackResponse(resp)
  1595  					if !bytes.Equal(data, []byte("OK")) {
  1596  						panic("invalid response")
  1597  					}
  1598  				}
  1599  				wg.Done()
  1600  			}()
  1601  			wg.Done()
  1602  		}()
  1603  	}
  1604  
  1605  	wg.Wait()
  1606  
  1607  	b.StopTimer()
  1608  	nsqd.Exit()
  1609  }
  1610  
  1611  func BenchmarkUnixSocketProtocolV2Pub256(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 256) }
  1612  func BenchmarkUnixSocketProtocolV2Pub512(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 512) }
  1613  func BenchmarkUnixSocketProtocolV2Pub1k(b *testing.B)  { benchmarkUnixSocketProtocolV2Pub(b, 1024) }
  1614  func BenchmarkUnixSocketProtocolV2Pub2k(b *testing.B)  { benchmarkUnixSocketProtocolV2Pub(b, 2*1024) }
  1615  func BenchmarkUnixSocketProtocolV2Pub4k(b *testing.B)  { benchmarkUnixSocketProtocolV2Pub(b, 4*1024) }
  1616  func BenchmarkUnixSocketProtocolV2Pub8k(b *testing.B)  { benchmarkUnixSocketProtocolV2Pub(b, 8*1024) }
  1617  func BenchmarkUnixSocketProtocolV2Pub16k(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 16*1024) }
  1618  func BenchmarkUnixSocketProtocolV2Pub32k(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 32*1024) }
  1619  func BenchmarkUnixSocketProtocolV2Pub64k(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 64*1024) }
  1620  func BenchmarkUnixSocketProtocolV2Pub128k(b *testing.B) {
  1621  	benchmarkUnixSocketProtocolV2Pub(b, 128*1024)
  1622  }
  1623  func BenchmarkUnixSocketProtocolV2Pub256k(b *testing.B) {
  1624  	benchmarkUnixSocketProtocolV2Pub(b, 256*1024)
  1625  }
  1626  func BenchmarkUnixSocketProtocolV2Pub512k(b *testing.B) {
  1627  	benchmarkUnixSocketProtocolV2Pub(b, 512*1024)
  1628  }
  1629  func BenchmarkUnixSocketProtocolV2Pub1m(b *testing.B) { benchmarkUnixSocketProtocolV2Pub(b, 1024*1024) }
  1630  
  1631  func benchmarkUnixSocketProtocolV2Sub(b *testing.B, size int) {
  1632  	var wg sync.WaitGroup
  1633  	b.StopTimer()
  1634  	opts := NewOptions()
  1635  	opts.Logger = test.NewTestLogger(b)
  1636  	opts.MemQueueSize = int64(b.N)
  1637  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1638  	defer os.RemoveAll(opts.DataPath)
  1639  	msg := make([]byte, size)
  1640  	topicName := "bench_v2_sub" + strconv.Itoa(b.N) + strconv.Itoa(int(time.Now().Unix()))
  1641  	topic := nsqd.GetTopic(topicName)
  1642  	for i := 0; i < b.N; i++ {
  1643  		msg := NewMessage(topic.GenerateID(), msg)
  1644  		topic.PutMessage(msg)
  1645  	}
  1646  	topic.GetChannel("ch")
  1647  	b.SetBytes(int64(len(msg)))
  1648  	goChan := make(chan int)
  1649  	rdyChan := make(chan int)
  1650  	workers := runtime.GOMAXPROCS(0)
  1651  	for j := 0; j < workers; j++ {
  1652  		wg.Add(1)
  1653  		go func() {
  1654  			subWorker(b.N, workers, addr, topicName, rdyChan, goChan)
  1655  			wg.Done()
  1656  		}()
  1657  		<-rdyChan
  1658  	}
  1659  	b.StartTimer()
  1660  
  1661  	close(goChan)
  1662  	wg.Wait()
  1663  
  1664  	b.StopTimer()
  1665  	nsqd.Exit()
  1666  }
  1667  
  1668  func subUnixSocketWorker(n int, workers int, addr net.Addr, topicName string, rdyChan chan int, goChan chan int) {
  1669  	conn, err := mustUnixSocketConnectNSQD(addr)
  1670  	if err != nil {
  1671  		panic(err.Error())
  1672  	}
  1673  	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriterSize(conn, 65536))
  1674  
  1675  	identify(nil, conn, nil, frameTypeResponse)
  1676  	sub(nil, conn, topicName, "ch")
  1677  
  1678  	rdyCount := int(math.Min(math.Max(float64(n/workers), 1), 2500))
  1679  	rdyChan <- 1
  1680  	<-goChan
  1681  	nsq.Ready(rdyCount).WriteTo(rw)
  1682  	rw.Flush()
  1683  	num := n / workers
  1684  	for i := 0; i < num; i++ {
  1685  		resp, err := nsq.ReadResponse(rw)
  1686  		if err != nil {
  1687  			panic(err.Error())
  1688  		}
  1689  		frameType, data, err := nsq.UnpackResponse(resp)
  1690  		if err != nil {
  1691  			panic(err.Error())
  1692  		}
  1693  		if frameType != frameTypeMessage {
  1694  			panic("got something else")
  1695  		}
  1696  		msg, err := decodeMessage(data)
  1697  		if err != nil {
  1698  			panic(err.Error())
  1699  		}
  1700  		nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(rw)
  1701  		if (i+1)%rdyCount == 0 || i+1 == num {
  1702  			if i+1 == num {
  1703  				nsq.Ready(0).WriteTo(conn)
  1704  			}
  1705  			rw.Flush()
  1706  		}
  1707  	}
  1708  }
  1709  
  1710  func BenchmarkUnixSocketProtocolV2Sub256(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 256) }
  1711  func BenchmarkUnixSocketProtocolV2Sub512(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 512) }
  1712  func BenchmarkUnixSocketProtocolV2Sub1k(b *testing.B)  { benchmarkUnixSocketProtocolV2Sub(b, 1024) }
  1713  func BenchmarkUnixSocketProtocolV2Sub2k(b *testing.B)  { benchmarkUnixSocketProtocolV2Sub(b, 2*1024) }
  1714  func BenchmarkUnixSocketProtocolV2Sub4k(b *testing.B)  { benchmarkUnixSocketProtocolV2Sub(b, 4*1024) }
  1715  func BenchmarkUnixSocketProtocolV2Sub8k(b *testing.B)  { benchmarkUnixSocketProtocolV2Sub(b, 8*1024) }
  1716  func BenchmarkUnixSocketProtocolV2Sub16k(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 16*1024) }
  1717  func BenchmarkUnixSocketProtocolV2Sub32k(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 32*1024) }
  1718  func BenchmarkUnixSocketProtocolV2Sub64k(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 64*1024) }
  1719  func BenchmarkUnixSocketProtocolV2Sub128k(b *testing.B) {
  1720  	benchmarkUnixSocketProtocolV2Sub(b, 128*1024)
  1721  }
  1722  func BenchmarkUnixSocketProtocolV2Sub256k(b *testing.B) {
  1723  	benchmarkUnixSocketProtocolV2Sub(b, 256*1024)
  1724  }
  1725  func BenchmarkUnixSocketProtocolV2Sub512k(b *testing.B) {
  1726  	benchmarkUnixSocketProtocolV2Sub(b, 512*1024)
  1727  }
  1728  func BenchmarkUnixSocketProtocolV2Sub1m(b *testing.B) { benchmarkUnixSocketProtocolV2Sub(b, 1024*1024) }
  1729  
  1730  func benchmarkUnixSocketProtocolV2MultiSub(b *testing.B, num int) {
  1731  	var wg sync.WaitGroup
  1732  	b.StopTimer()
  1733  
  1734  	opts := NewOptions()
  1735  	opts.Logger = test.NewTestLogger(b)
  1736  	opts.MemQueueSize = int64(b.N)
  1737  	addr, _, nsqd := mustUnixSocketStartNSQD(opts)
  1738  	defer os.RemoveAll(opts.DataPath)
  1739  	msg := make([]byte, 256)
  1740  	b.SetBytes(int64(len(msg) * num))
  1741  
  1742  	goChan := make(chan int)
  1743  	rdyChan := make(chan int)
  1744  	workers := runtime.GOMAXPROCS(0)
  1745  	for i := 0; i < num; i++ {
  1746  		topicName := "bench_v2" + strconv.Itoa(b.N) + "_" + strconv.Itoa(i) + "_" + strconv.Itoa(int(time.Now().Unix()))
  1747  		topic := nsqd.GetTopic(topicName)
  1748  		for i := 0; i < b.N; i++ {
  1749  			msg := NewMessage(topic.GenerateID(), msg)
  1750  			topic.PutMessage(msg)
  1751  		}
  1752  		topic.GetChannel("ch")
  1753  
  1754  		for j := 0; j < workers; j++ {
  1755  			wg.Add(1)
  1756  			go func() {
  1757  				subUnixSocketWorker(b.N, workers, addr, topicName, rdyChan, goChan)
  1758  				wg.Done()
  1759  			}()
  1760  			<-rdyChan
  1761  		}
  1762  	}
  1763  	b.StartTimer()
  1764  
  1765  	close(goChan)
  1766  	wg.Wait()
  1767  
  1768  	b.StopTimer()
  1769  	nsqd.Exit()
  1770  }
  1771  
  1772  func BenchmarkUnixSocketProtocolV2MultiSub2(b *testing.B) {
  1773  	benchmarkUnixSocketProtocolV2MultiSub(b, 2)
  1774  }
  1775  func BenchmarkUnixSocketProtocolV2MultiSub1(b *testing.B) {
  1776  	benchmarkUnixSocketProtocolV2MultiSub(b, 1)
  1777  }
  1778  func BenchmarkUnixSocketProtocolV2MultiSub4(b *testing.B) {
  1779  	benchmarkUnixSocketProtocolV2MultiSub(b, 4)
  1780  }
  1781  func BenchmarkUnixSocketProtocolV2MultiSub8(b *testing.B) {
  1782  	benchmarkUnixSocketProtocolV2MultiSub(b, 8)
  1783  }
  1784  func BenchmarkUnixSocketProtocolV2MultiSub16(b *testing.B) {
  1785  	benchmarkUnixSocketProtocolV2MultiSub(b, 16)
  1786  }