func testIOLoopReturnsClientErr(t *testing.T, fakeConn test.FakeNetConn) {
fakeConn.ReadFunc = func(b []byte) (int, error) {
return copy(b, []byte("INVALID_COMMAND\n")), nil
}
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
prot := &LookupProtocolV1{ctx: &Context{nsqlookupd: New(opts)}}
errChan := make(chan error)
testIOLoop := func() {
errChan <- prot.IOLoop(fakeConn)
defer prot.ctx.nsqlookupd.Exit()
}
go testIOLoop()
var err error
var timeout bool
select {
case err = <-errChan:
case <-time.After(2 * time.Second):
timeout = true
}
test.Equal(t, false, timeout)
test.NotNil(t, err)
test.Equal(t, "E_INVALID invalid command INVALID_COMMAND", err.Error())
test.NotNil(t, err.(*protocol.FatalClientErr))
}
func TestChannelHealth(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.MemQueueSize = 2
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topic := nsqd.GetTopic("test")
channel := topic.GetChannel("channel")
channel.backend = &errorBackendQueue{}
msg := NewMessage(<-nsqd.idChan, make([]byte, 100))
err := channel.PutMessage(msg)
test.Nil(t, err)
msg = NewMessage(<-nsqd.idChan, make([]byte, 100))
err = channel.PutMessage(msg)
test.Nil(t, err)
msg = NewMessage(<-nsqd.idChan, make([]byte, 100))
err = channel.PutMessage(msg)
test.NotNil(t, err)
url := fmt.Sprintf("http://%s/ping", httpAddr)
resp, err := http.Get(url)
test.Nil(t, err)
test.Equal(t, 500, resp.StatusCode)
body, _ := ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "NOK - never gonna happen", string(body))
channel.backend = &errorRecoveredBackendQueue{}
msg = NewMessage(<-nsqd.idChan, make([]byte, 100))
err = channel.PutMessage(msg)
test.Nil(t, err)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "OK", string(body))
}
func TestHTTPgetStatusText(t *testing.T) {
testTime := time.Now()
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
nsqd.startTime = testTime
url := fmt.Sprintf("http://%s/stats?format=text", httpAddr)
resp, err := http.Get(url)
test.Nil(t, err)
defer resp.Body.Close()
body, _ := ioutil.ReadAll(resp.Body)
test.Equal(t, 200, resp.StatusCode)
test.NotNil(t, body)
}
func TestDiskQueue(t *testing.T) {
l := test.NewTestLogger(t)
dqName := "test_disk_queue" + strconv.Itoa(int(time.Now().Unix()))
tmpDir, err := ioutil.TempDir("", fmt.Sprintf("nsq-test-%d", time.Now().UnixNano()))
if err != nil {
panic(err)
}
defer os.RemoveAll(tmpDir)
dq := newDiskQueue(dqName, tmpDir, 1024, 4, 1<<10, 2500, 2*time.Second, l)
defer dq.Close()
test.NotNil(t, dq)
test.Equal(t, int64(0), dq.Depth())
msg := []byte("test")
err = dq.Put(msg)
test.Nil(t, err)
test.Equal(t, int64(1), dq.Depth())
msgOut := <-dq.ReadChan()
test.Equal(t, msg, msgOut)
}
func TestSetHealth(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
nsqd := New(opts)
test.Equal(t, nil, nsqd.GetError())
test.Equal(t, true, nsqd.IsHealthy())
nsqd.SetHealth(nil)
test.Equal(t, nil, nsqd.GetError())
test.Equal(t, true, nsqd.IsHealthy())
nsqd.SetHealth(errors.New("health error"))
test.NotNil(t, nsqd.GetError())
test.Equal(t, "NOK - health error", nsqd.GetHealth())
test.Equal(t, false, nsqd.IsHealthy())
nsqd.SetHealth(nil)
test.Nil(t, nsqd.GetError())
test.Equal(t, "OK", nsqd.GetHealth())
test.Equal(t, true, nsqd.IsHealthy())
}
func TestDiskQueueRoll(t *testing.T) {
l := test.NewTestLogger(t)
dqName := "test_disk_queue_roll" + strconv.Itoa(int(time.Now().Unix()))
tmpDir, err := ioutil.TempDir("", fmt.Sprintf("nsq-test-%d", time.Now().UnixNano()))
if err != nil {
panic(err)
}
defer os.RemoveAll(tmpDir)
msg := bytes.Repeat([]byte{0}, 10)
ml := int64(len(msg))
dq := newDiskQueue(dqName, tmpDir, 9*(ml+4), int32(ml), 1<<10, 2500, 2*time.Second, l)
defer dq.Close()
test.NotNil(t, dq)
test.Equal(t, int64(0), dq.Depth())
for i := 0; i < 10; i++ {
err := dq.Put(msg)
test.Nil(t, err)
test.Equal(t, int64(i+1), dq.Depth())
}
test.Equal(t, int64(1), dq.(*diskQueue).writeFileNum)
test.Equal(t, int64(0), dq.(*diskQueue).writePos)
}
func TestHTTPV1TopicChannel(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_http_topic_channel2" + strconv.Itoa(int(time.Now().Unix()))
channelName := "ch2"
url := fmt.Sprintf("http://%s/topic/create?topic=%s", httpAddr, topicName)
resp, err := http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ := ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
url = fmt.Sprintf("http://%s/channel/create?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
topic, err := nsqd.GetExistingTopic(topicName)
test.Nil(t, err)
test.NotNil(t, topic)
channel, err := topic.GetExistingChannel(channelName)
test.Nil(t, err)
test.NotNil(t, channel)
em := ErrMessage{}
url = fmt.Sprintf("http://%s/topic/pause", httpAddr)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 400, resp.StatusCode)
test.Equal(t, "Bad Request", http.StatusText(resp.StatusCode))
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
t.Logf("%s", body)
err = json.Unmarshal(body, &em)
test.Nil(t, err)
test.Equal(t, "MISSING_ARG_TOPIC", em.Message)
url = fmt.Sprintf("http://%s/topic/pause?topic=%s", httpAddr, topicName+"abc")
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 404, resp.StatusCode)
test.Equal(t, "Not Found", http.StatusText(resp.StatusCode))
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
t.Logf("%s", body)
err = json.Unmarshal(body, &em)
test.Nil(t, err)
test.Equal(t, "TOPIC_NOT_FOUND", em.Message)
url = fmt.Sprintf("http://%s/topic/pause?topic=%s", httpAddr, topicName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, true, topic.IsPaused())
url = fmt.Sprintf("http://%s/topic/unpause?topic=%s", httpAddr, topicName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, false, topic.IsPaused())
url = fmt.Sprintf("http://%s/channel/pause?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, true, channel.IsPaused())
url = fmt.Sprintf("http://%s/channel/unpause?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, false, channel.IsPaused())
url = fmt.Sprintf("http://%s/channel/delete?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
_, err = topic.GetExistingChannel(channelName)
test.NotNil(t, err)
url = fmt.Sprintf("http://%s/topic/delete?topic=%s", httpAddr, topicName)
resp, err = http.Post(url, "application/json", nil)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
_, err = nsqd.GetExistingTopic(topicName)
test.NotNil(t, err)
}
func TestHTTPTransitionTopicChannel(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
client := http.Client{}
topicName := "test_http_topic_channel1" + strconv.Itoa(int(time.Now().Unix()))
channelName := "ch1"
url := fmt.Sprintf("http://%s/create_topic?topic=%s", httpAddr, topicName)
req, _ := http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err := client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ := ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
url = fmt.Sprintf("http://%s/create_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
topic, err := nsqd.GetExistingTopic(topicName)
test.Nil(t, err)
test.NotNil(t, topic)
channel, err := topic.GetExistingChannel(channelName)
test.Nil(t, err)
test.NotNil(t, channel)
url = fmt.Sprintf("http://%s/pause_topic?topic=%s", httpAddr, topicName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, true, topic.IsPaused())
url = fmt.Sprintf("http://%s/unpause_topic?topic=%s", httpAddr, topicName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, false, topic.IsPaused())
url = fmt.Sprintf("http://%s/pause_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, true, channel.IsPaused())
url = fmt.Sprintf("http://%s/unpause_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
test.Equal(t, false, channel.IsPaused())
url = fmt.Sprintf("http://%s/delete_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
_, err = topic.GetExistingChannel(channelName)
test.NotNil(t, err)
url = fmt.Sprintf("http://%s/delete_topic?topic=%s", httpAddr, topicName)
req, _ = http.NewRequest("GET", url, nil)
req.Header.Set("Accept", "application/vnd.nsq; version=1.0")
resp, err = client.Do(req)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, "", string(body))
test.Equal(t, "nsq; version=1.0", resp.Header.Get("X-NSQ-Content-Type"))
_, err = nsqd.GetExistingTopic(topicName)
test.NotNil(t, err)
}
func TestHTTPDeprecatedTopicChannel(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_http_topic_channel" + strconv.Itoa(int(time.Now().Unix()))
channelName := "ch"
url := fmt.Sprintf("http://%s/create_topic?topic=%s", httpAddr, topicName)
resp, err := http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ := ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
url = fmt.Sprintf("http://%s/create_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
topic, err := nsqd.GetExistingTopic(topicName)
test.Nil(t, err)
test.NotNil(t, topic)
channel, err := topic.GetExistingChannel(channelName)
test.Nil(t, err)
test.NotNil(t, channel)
url = fmt.Sprintf("http://%s/pause_topic?topic=%s", httpAddr, topicName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
test.Equal(t, true, topic.IsPaused())
url = fmt.Sprintf("http://%s/unpause_topic?topic=%s", httpAddr, topicName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
test.Equal(t, false, topic.IsPaused())
url = fmt.Sprintf("http://%s/pause_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
test.Equal(t, true, channel.IsPaused())
url = fmt.Sprintf("http://%s/unpause_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
test.Equal(t, false, channel.IsPaused())
url = fmt.Sprintf("http://%s/delete_channel?topic=%s&channel=%s", httpAddr, topicName, channelName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
_, err = topic.GetExistingChannel(channelName)
test.NotNil(t, err)
url = fmt.Sprintf("http://%s/delete_topic?topic=%s", httpAddr, topicName)
resp, err = http.Get(url)
test.Nil(t, err)
test.Equal(t, 200, resp.StatusCode)
body, _ = ioutil.ReadAll(resp.Body)
resp.Body.Close()
test.Equal(t, `{"status_code":200,"status_txt":"OK","data":null}`, string(body))
_, err = nsqd.GetExistingTopic(topicName)
test.NotNil(t, err)
}
func TestHTTPSRequireVerify(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.TLSCert = "./test/certs/server.pem"
opts.TLSKey = "./test/certs/server.key"
opts.TLSRootCAFile = "./test/certs/ca.pem"
opts.TLSClientAuthPolicy = "require-verify"
_, httpAddr, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
httpsAddr := nsqd.httpsListener.Addr().(*net.TCPAddr)
topicName := "test_http_put_req_verf" + strconv.Itoa(int(time.Now().Unix()))
topic := nsqd.GetTopic(topicName)
// no cert
buf := bytes.NewBuffer([]byte("test message"))
url := fmt.Sprintf("http://%s/put?topic=%s", httpAddr, topicName)
resp, err := http.Post(url, "application/octet-stream", buf)
test.Equal(t, 403, resp.StatusCode)
// unsigned cert
cert, err := tls.LoadX509KeyPair("./test/certs/cert.pem", "./test/certs/key.pem")
test.Nil(t, err)
tlsConfig := &tls.Config{
Certificates: []tls.Certificate{cert},
InsecureSkipVerify: true,
}
transport := &http.Transport{
TLSClientConfig: tlsConfig,
}
client := &http.Client{Transport: transport}
buf = bytes.NewBuffer([]byte("test message"))
url = fmt.Sprintf("https://%s/put?topic=%s", httpsAddr, topicName)
resp, err = client.Post(url, "application/octet-stream", buf)
test.NotNil(t, err)
// signed cert
cert, err = tls.LoadX509KeyPair("./test/certs/client.pem", "./test/certs/client.key")
test.Nil(t, err)
tlsConfig = &tls.Config{
Certificates: []tls.Certificate{cert},
InsecureSkipVerify: true,
}
transport = &http.Transport{
TLSClientConfig: tlsConfig,
}
client = &http.Client{Transport: transport}
buf = bytes.NewBuffer([]byte("test message"))
url = fmt.Sprintf("https://%s/put?topic=%s", httpsAddr, topicName)
resp, err = client.Post(url, "application/octet-stream", buf)
test.Nil(t, err)
defer resp.Body.Close()
body, _ := ioutil.ReadAll(resp.Body)
test.Equal(t, "OK", string(body))
time.Sleep(5 * time.Millisecond)
test.Equal(t, int64(1), topic.Depth())
}
func TestDiskQueueEmpty(t *testing.T) {
l := test.NewTestLogger(t)
dqName := "test_disk_queue_empty" + strconv.Itoa(int(time.Now().Unix()))
tmpDir, err := ioutil.TempDir("", fmt.Sprintf("nsq-test-%d", time.Now().UnixNano()))
if err != nil {
panic(err)
}
defer os.RemoveAll(tmpDir)
msg := bytes.Repeat([]byte{0}, 10)
dq := newDiskQueue(dqName, tmpDir, 100, 0, 1<<10, 2500, 2*time.Second, l)
defer dq.Close()
test.NotNil(t, dq)
test.Equal(t, int64(0), dq.Depth())
for i := 0; i < 100; i++ {
err := dq.Put(msg)
test.Nil(t, err)
test.Equal(t, int64(i+1), dq.Depth())
}
for i := 0; i < 3; i++ {
<-dq.ReadChan()
}
for {
if dq.Depth() == 97 {
break
}
time.Sleep(50 * time.Millisecond)
}
test.Equal(t, int64(97), dq.Depth())
numFiles := dq.(*diskQueue).writeFileNum
dq.Empty()
assertFileNotExist(t, dq.(*diskQueue).metaDataFileName())
for i := int64(0); i <= numFiles; i++ {
assertFileNotExist(t, dq.(*diskQueue).fileName(i))
}
test.Equal(t, int64(0), dq.Depth())
test.Equal(t, dq.(*diskQueue).writeFileNum, dq.(*diskQueue).readFileNum)
test.Equal(t, dq.(*diskQueue).writePos, dq.(*diskQueue).readPos)
test.Equal(t, dq.(*diskQueue).readPos, dq.(*diskQueue).nextReadPos)
test.Equal(t, dq.(*diskQueue).readFileNum, dq.(*diskQueue).nextReadFileNum)
for i := 0; i < 100; i++ {
err := dq.Put(msg)
test.Nil(t, err)
test.Equal(t, int64(i+1), dq.Depth())
}
for i := 0; i < 100; i++ {
<-dq.ReadChan()
}
for {
if dq.Depth() == 0 {
break
}
time.Sleep(50 * time.Millisecond)
}
test.Equal(t, int64(0), dq.Depth())
test.Equal(t, dq.(*diskQueue).writeFileNum, dq.(*diskQueue).readFileNum)
test.Equal(t, dq.(*diskQueue).writePos, dq.(*diskQueue).readPos)
test.Equal(t, dq.(*diskQueue).readPos, dq.(*diskQueue).nextReadPos)
}
func TestDiskQueueTorture(t *testing.T) {
var wg sync.WaitGroup
l := test.NewTestLogger(t)
dqName := "test_disk_queue_torture" + strconv.Itoa(int(time.Now().Unix()))
tmpDir, err := ioutil.TempDir("", fmt.Sprintf("nsq-test-%d", time.Now().UnixNano()))
if err != nil {
panic(err)
}
defer os.RemoveAll(tmpDir)
dq := newDiskQueue(dqName, tmpDir, 262144, 0, 1<<10, 2500, 2*time.Second, l)
test.NotNil(t, dq)
test.Equal(t, int64(0), dq.Depth())
msg := []byte("aaaaaaaaaabbbbbbbbbbccccccccccddddddddddeeeeeeeeeeffffffffff")
numWriters := 4
numReaders := 4
readExitChan := make(chan int)
writeExitChan := make(chan int)
var depth int64
for i := 0; i < numWriters; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for {
time.Sleep(100000 * time.Nanosecond)
select {
case <-writeExitChan:
return
default:
err := dq.Put(msg)
if err == nil {
atomic.AddInt64(&depth, 1)
}
}
}
}()
}
time.Sleep(1 * time.Second)
dq.Close()
t.Logf("closing writeExitChan")
close(writeExitChan)
wg.Wait()
t.Logf("restarting diskqueue")
dq = newDiskQueue(dqName, tmpDir, 262144, 0, 1<<10, 2500, 2*time.Second, l)
defer dq.Close()
test.NotNil(t, dq)
test.Equal(t, depth, dq.Depth())
var read int64
for i := 0; i < numReaders; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for {
time.Sleep(100000 * time.Nanosecond)
select {
case m := <-dq.ReadChan():
test.Equal(t, m, msg)
atomic.AddInt64(&read, 1)
case <-readExitChan:
return
}
}
}()
}
t.Logf("waiting for depth 0")
for {
if dq.Depth() == 0 {
break
}
time.Sleep(50 * time.Millisecond)
}
t.Logf("closing readExitChan")
close(readExitChan)
wg.Wait()
test.Equal(t, depth, read)
}
func TestClientTimeout(t *testing.T) {
topicName := "test_client_timeout_v2" + strconv.Itoa(int(time.Now().Unix()))
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.ClientTimeout = 150 * time.Millisecond
opts.Verbose = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
time.Sleep(150 * time.Millisecond)
// depending on timing there may be 1 or 2 hearbeats sent
// just read until we get an error
timer := time.After(100 * time.Millisecond)
for {
select {
case <-timer:
t.Fatalf("test timed out")
default:
_, err := nsq.ReadResponse(conn)
if err != nil {
goto done
}
}
}
done:
}
func TestClientHeartbeat(t *testing.T) {
topicName := "test_hb_v2" + strconv.Itoa(int(time.Now().Unix()))
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.ClientTimeout = 200 * time.Millisecond
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(1).WriteTo(conn)
test.Nil(t, err)
resp, _ := nsq.ReadResponse(conn)
_, data, _ := nsq.UnpackResponse(resp)
test.Equal(t, []byte("_heartbeat_"), data)
time.Sleep(20 * time.Millisecond)
_, err = nsq.Nop().WriteTo(conn)
test.Nil(t, err)
// wait long enough that would have timed out (had we not sent the above cmd)
time.Sleep(100 * time.Millisecond)
_, err = nsq.Nop().WriteTo(conn)
test.Nil(t, err)
}
func TestClientHeartbeatDisableSUB(t *testing.T) {
topicName := "test_hb_v2" + strconv.Itoa(int(time.Now().Unix()))
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.ClientTimeout = 200 * time.Millisecond
opts.Verbose = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, map[string]interface{}{
"heartbeat_interval": -1,
}, frameTypeResponse)
subFail(t, conn, topicName, "ch")
}
func TestClientHeartbeatDisable(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.ClientTimeout = 100 * time.Millisecond
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, map[string]interface{}{
"heartbeat_interval": -1,
}, frameTypeResponse)
time.Sleep(150 * time.Millisecond)
_, err = nsq.Nop().WriteTo(conn)
test.Nil(t, err)
}
func TestMaxHeartbeatIntervalValid(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.MaxHeartbeatInterval = 300 * time.Second
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
hbi := int(opts.MaxHeartbeatInterval / time.Millisecond)
identify(t, conn, map[string]interface{}{
"heartbeat_interval": hbi,
}, frameTypeResponse)
}
func TestMaxHeartbeatIntervalInvalid(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.MaxHeartbeatInterval = 300 * time.Second
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
hbi := int(opts.MaxHeartbeatInterval/time.Millisecond + 1)
data := identify(t, conn, map[string]interface{}{
"heartbeat_interval": hbi,
}, frameTypeError)
test.Equal(t, "E_BAD_BODY IDENTIFY heartbeat interval (300001) is invalid", string(data))
}
func TestPausing(t *testing.T) {
topicName := "test_pause_v2" + strconv.Itoa(int(time.Now().Unix()))
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(1).WriteTo(conn)
test.Nil(t, err)
topic := nsqd.GetTopic(topicName)
msg := NewMessage(<-nsqd.idChan, []byte("test body"))
channel := topic.GetChannel("ch")
topic.PutMessage(msg)
// receive the first message via the client, finish it, and send new RDY
resp, _ := nsq.ReadResponse(conn)
_, data, _ := nsq.UnpackResponse(resp)
msg, err = decodeMessage(data)
test.Equal(t, []byte("test body"), msg.Body)
_, err = nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(conn)
test.Nil(t, err)
_, err = nsq.Ready(1).WriteTo(conn)
test.Nil(t, err)
// sleep to allow the RDY state to take effect
time.Sleep(50 * time.Millisecond)
// pause the channel... the client shouldn't receive any more messages
channel.Pause()
// sleep to allow the paused state to take effect
time.Sleep(50 * time.Millisecond)
msg = NewMessage(<-nsqd.idChan, []byte("test body2"))
topic.PutMessage(msg)
// allow the client to possibly get a message, the test would hang indefinitely
// if pausing was not working
time.Sleep(50 * time.Millisecond)
msg = <-channel.memoryMsgChan
test.Equal(t, []byte("test body2"), msg.Body)
// unpause the channel... the client should now be pushed a message
channel.UnPause()
msg = NewMessage(<-nsqd.idChan, []byte("test body3"))
topic.PutMessage(msg)
resp, _ = nsq.ReadResponse(conn)
_, data, _ = nsq.UnpackResponse(resp)
msg, err = decodeMessage(data)
test.Equal(t, []byte("test body3"), msg.Body)
}
func TestEmptyCommand(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
_, err = conn.Write([]byte("\n\n"))
test.Nil(t, err)
// if we didn't panic here we're good, see issue #120
}
func TestSizeLimits(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MaxMsgSize = 100
opts.MaxBodySize = 1000
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
topicName := "test_limits_v2" + strconv.Itoa(int(time.Now().Unix()))
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
// PUB that's valid
nsq.Publish(topicName, make([]byte, 95)).WriteTo(conn)
resp, _ := nsq.ReadResponse(conn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
// PUB that's invalid (too big)
nsq.Publish(topicName, make([]byte, 105)).WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_BAD_MESSAGE PUB message too big 105 > 100"), string(data))
// need to reconnect
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
// PUB thats empty
nsq.Publish(topicName, []byte{}).WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_BAD_MESSAGE PUB invalid message body size 0"), string(data))
// need to reconnect
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
// MPUB body that's valid
mpub := make([][]byte, 5)
for i := range mpub {
mpub[i] = make([]byte, 100)
}
cmd, _ := nsq.MultiPublish(topicName, mpub)
cmd.WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
// MPUB body that's invalid (body too big)
mpub = make([][]byte, 11)
for i := range mpub {
mpub[i] = make([]byte, 100)
}
cmd, _ = nsq.MultiPublish(topicName, mpub)
cmd.WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_BAD_BODY MPUB body too big 1148 > 1000"), string(data))
// need to reconnect
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
// MPUB that's invalid (one message empty)
mpub = make([][]byte, 5)
for i := range mpub {
mpub[i] = make([]byte, 100)
}
mpub = append(mpub, []byte{})
cmd, _ = nsq.MultiPublish(topicName, mpub)
cmd.WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_BAD_MESSAGE MPUB invalid message(5) body size 0"), string(data))
// need to reconnect
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
// MPUB body that's invalid (one of the messages is too big)
mpub = make([][]byte, 5)
for i := range mpub {
mpub[i] = make([]byte, 101)
}
cmd, _ = nsq.MultiPublish(topicName, mpub)
cmd.WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_BAD_MESSAGE MPUB message too big 101 > 100"), string(data))
}
func TestDPUB(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
topicName := "test_dpub_v2" + strconv.Itoa(int(time.Now().Unix()))
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
// valid
nsq.DeferredPublish(topicName, time.Second, make([]byte, 100)).WriteTo(conn)
resp, _ := nsq.ReadResponse(conn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
time.Sleep(25 * time.Millisecond)
ch := nsqd.GetTopic(topicName).GetChannel("ch")
ch.deferredMutex.Lock()
numDef := len(ch.deferredMessages)
ch.deferredMutex.Unlock()
test.Equal(t, 1, numDef)
// duration out of range
nsq.DeferredPublish(topicName, opts.MaxReqTimeout+100*time.Millisecond, make([]byte, 100)).WriteTo(conn)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_INVALID DPUB timeout 3600100 out of range 0-3600000"), string(data))
}
func TestTouch(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MsgTimeout = 150 * time.Millisecond
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_touch" + strconv.Itoa(int(time.Now().Unix()))
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, nil, frameTypeResponse)
sub(t, conn, topicName, "ch")
topic := nsqd.GetTopic(topicName)
channel := topic.GetChannel("ch")
msg := NewMessage(<-nsqd.idChan, []byte("test body"))
topic.PutMessage(msg)
_, err = nsq.Ready(1).WriteTo(conn)
test.Nil(t, err)
resp, err := nsq.ReadResponse(conn)
test.Nil(t, err)
frameType, data, err := nsq.UnpackResponse(resp)
msgOut, _ := decodeMessage(data)
test.Equal(t, frameTypeMessage, frameType)
test.Equal(t, msg.ID, msgOut.ID)
time.Sleep(75 * time.Millisecond)
_, err = nsq.Touch(nsq.MessageID(msg.ID)).WriteTo(conn)
test.Nil(t, err)
time.Sleep(75 * time.Millisecond)
_, err = nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(conn)
test.Nil(t, err)
test.Equal(t, uint64(0), channel.timeoutCount)
}
func TestMaxRdyCount(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MaxRdyCount = 50
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_max_rdy_count" + strconv.Itoa(int(time.Now().Unix()))
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
topic := nsqd.GetTopic(topicName)
msg := NewMessage(<-nsqd.idChan, []byte("test body"))
topic.PutMessage(msg)
data := identify(t, conn, nil, frameTypeResponse)
r := struct {
MaxRdyCount int64 `json:"max_rdy_count"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, int64(50), r.MaxRdyCount)
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(int(opts.MaxRdyCount)).WriteTo(conn)
test.Nil(t, err)
resp, err := nsq.ReadResponse(conn)
test.Nil(t, err)
frameType, data, err := nsq.UnpackResponse(resp)
msgOut, _ := decodeMessage(data)
test.Equal(t, frameTypeMessage, frameType)
test.Equal(t, msg.ID, msgOut.ID)
_, err = nsq.Ready(int(opts.MaxRdyCount) + 1).WriteTo(conn)
test.Nil(t, err)
resp, err = nsq.ReadResponse(conn)
test.Nil(t, err)
frameType, data, err = nsq.UnpackResponse(resp)
test.Equal(t, int32(1), frameType)
test.Equal(t, "E_INVALID RDY count 51 out of range 0-50", string(data))
}
func TestFatalError(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
_, err = conn.Write([]byte("ASDF\n"))
test.Nil(t, err)
resp, err := nsq.ReadResponse(conn)
test.Nil(t, err)
frameType, data, err := nsq.UnpackResponse(resp)
test.Equal(t, int32(1), frameType)
test.Equal(t, "E_INVALID invalid command ASDF", string(data))
_, err = nsq.ReadResponse(conn)
test.NotNil(t, err)
}
func TestOutputBuffering(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MaxOutputBufferSize = 512 * 1024
opts.MaxOutputBufferTimeout = time.Second
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_output_buffering" + strconv.Itoa(int(time.Now().Unix()))
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
outputBufferSize := 256 * 1024
outputBufferTimeout := 500
topic := nsqd.GetTopic(topicName)
msg := NewMessage(<-nsqd.idChan, make([]byte, outputBufferSize-1024))
topic.PutMessage(msg)
start := time.Now()
data := identify(t, conn, map[string]interface{}{
"output_buffer_size": outputBufferSize,
"output_buffer_timeout": outputBufferTimeout,
}, frameTypeResponse)
var decoded map[string]interface{}
json.Unmarshal(data, &decoded)
v, ok := decoded["output_buffer_size"]
test.Equal(t, true, ok)
test.Equal(t, outputBufferSize, int(v.(float64)))
v, ok = decoded["output_buffer_timeout"]
test.Equal(t, outputBufferTimeout, int(v.(float64)))
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(10).WriteTo(conn)
test.Nil(t, err)
resp, err := nsq.ReadResponse(conn)
test.Nil(t, err)
end := time.Now()
test.Equal(t, true, int(end.Sub(start)/time.Millisecond) >= outputBufferTimeout)
frameType, data, err := nsq.UnpackResponse(resp)
msgOut, _ := decodeMessage(data)
test.Equal(t, frameTypeMessage, frameType)
test.Equal(t, msg.ID, msgOut.ID)
}
func TestOutputBufferingValidity(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MaxOutputBufferSize = 512 * 1024
opts.MaxOutputBufferTimeout = time.Second
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, map[string]interface{}{
"output_buffer_size": 512 * 1024,
"output_buffer_timeout": 1000,
}, frameTypeResponse)
identify(t, conn, map[string]interface{}{
"output_buffer_size": -1,
"output_buffer_timeout": -1,
}, frameTypeResponse)
identify(t, conn, map[string]interface{}{
"output_buffer_size": 0,
"output_buffer_timeout": 0,
}, frameTypeResponse)
data := identify(t, conn, map[string]interface{}{
"output_buffer_size": 512*1024 + 1,
"output_buffer_timeout": 0,
}, frameTypeError)
test.Equal(t, fmt.Sprintf("E_BAD_BODY IDENTIFY output buffer size (%d) is invalid", 512*1024+1), string(data))
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data = identify(t, conn, map[string]interface{}{
"output_buffer_size": 0,
"output_buffer_timeout": 1001,
}, frameTypeError)
test.Equal(t, "E_BAD_BODY IDENTIFY output buffer timeout (1001) is invalid", string(data))
}
func TestTLS(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.TLSCert = "./test/certs/server.pem"
opts.TLSKey = "./test/certs/server.key"
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestTLSRequired(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.TLSCert = "./test/certs/server.pem"
opts.TLSKey = "./test/certs/server.key"
opts.TLSRequired = TLSRequiredExceptHTTP
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_tls_required" + strconv.Itoa(int(time.Now().Unix()))
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
subFail(t, conn, topicName, "ch")
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestTLSAuthRequire(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.TLSCert = "./test/certs/server.pem"
opts.TLSKey = "./test/certs/server.key"
opts.TLSClientAuthPolicy = "require"
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
// No Certs
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.NotNil(t, err)
// With Unsigned Cert
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data = identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r = struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
cert, err := tls.LoadX509KeyPair("./test/certs/cert.pem", "./test/certs/key.pem")
test.Nil(t, err)
tlsConfig = &tls.Config{
Certificates: []tls.Certificate{cert},
InsecureSkipVerify: true,
}
tlsConn = tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestTLSAuthRequireVerify(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.TLSCert = "./test/certs/server.pem"
opts.TLSKey = "./test/certs/server.key"
opts.TLSRootCAFile = "./test/certs/ca.pem"
opts.TLSClientAuthPolicy = "require-verify"
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
// with no cert
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.NotNil(t, err)
// with invalid cert
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data = identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r = struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
cert, err := tls.LoadX509KeyPair("./test/certs/cert.pem", "./test/certs/key.pem")
test.Nil(t, err)
tlsConfig = &tls.Config{
Certificates: []tls.Certificate{cert},
InsecureSkipVerify: true,
}
tlsConn = tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.NotNil(t, err)
// with valid cert
conn, err = mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data = identify(t, conn, map[string]interface{}{
"tls_v1": true,
}, frameTypeResponse)
r = struct {
TLSv1 bool `json:"tls_v1"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
cert, err = tls.LoadX509KeyPair("./test/certs/client.pem", "./test/certs/client.key")
test.Nil(t, err)
tlsConfig = &tls.Config{
Certificates: []tls.Certificate{cert},
InsecureSkipVerify: true,
}
tlsConn = tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestDeflate(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.DeflateEnabled = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"deflate": true,
}, frameTypeResponse)
r := struct {
Deflate bool `json:"deflate"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.Deflate)
compressConn := flate.NewReader(conn)
resp, _ := nsq.ReadResponse(compressConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
type readWriter struct {
io.Reader
io.Writer
}
func TestSnappy(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.SnappyEnabled = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"snappy": true,
}, frameTypeResponse)
r := struct {
Snappy bool `json:"snappy"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.Snappy)
compressConn := snappystream.NewReader(conn, snappystream.SkipVerifyChecksum)
resp, _ := nsq.ReadResponse(compressConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
msgBody := make([]byte, 128000)
w := snappystream.NewWriter(conn)
rw := readWriter{compressConn, w}
topicName := "test_snappy" + strconv.Itoa(int(time.Now().Unix()))
sub(t, rw, topicName, "ch")
_, err = nsq.Ready(1).WriteTo(rw)
test.Nil(t, err)
topic := nsqd.GetTopic(topicName)
msg := NewMessage(<-nsqd.idChan, msgBody)
topic.PutMessage(msg)
resp, _ = nsq.ReadResponse(compressConn)
frameType, data, _ = nsq.UnpackResponse(resp)
msgOut, _ := decodeMessage(data)
test.Equal(t, frameTypeMessage, frameType)
test.Equal(t, msg.ID, msgOut.ID)
test.Equal(t, msg.Body, msgOut.Body)
}
func TestTLSDeflate(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.DeflateEnabled = true
opts.TLSCert = "./test/certs/cert.pem"
opts.TLSKey = "./test/certs/key.pem"
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
"deflate": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
Deflate bool `json:"deflate"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
test.Equal(t, true, r.Deflate)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
compressConn := flate.NewReader(tlsConn)
resp, _ = nsq.ReadResponse(compressConn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestSampling(t *testing.T) {
rand.Seed(time.Now().UTC().UnixNano())
num := 10000
sampleRate := 42
slack := 5
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.MaxRdyCount = int64(num)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"sample_rate": int32(sampleRate),
}, frameTypeResponse)
r := struct {
SampleRate int32 `json:"sample_rate"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, int32(sampleRate), r.SampleRate)
topicName := "test_sampling" + strconv.Itoa(int(time.Now().Unix()))
topic := nsqd.GetTopic(topicName)
for i := 0; i < num; i++ {
msg := NewMessage(<-nsqd.idChan, []byte("test body"))
topic.PutMessage(msg)
}
channel := topic.GetChannel("ch")
// let the topic drain into the channel
time.Sleep(50 * time.Millisecond)
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(num).WriteTo(conn)
test.Nil(t, err)
go func() {
for {
_, err := nsq.ReadResponse(conn)
if err != nil {
return
}
}
}()
doneChan := make(chan int)
go func() {
for {
if channel.Depth() == 0 {
close(doneChan)
return
}
time.Sleep(5 * time.Millisecond)
}
}()
<-doneChan
channel.inFlightMutex.Lock()
numInFlight := len(channel.inFlightMessages)
channel.inFlightMutex.Unlock()
test.Equal(t, true, numInFlight <= int(float64(num)*float64(sampleRate+slack)/100.0))
test.Equal(t, true, numInFlight >= int(float64(num)*float64(sampleRate-slack)/100.0))
}
func TestTLSSnappy(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.SnappyEnabled = true
opts.TLSCert = "./test/certs/cert.pem"
opts.TLSKey = "./test/certs/key.pem"
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
data := identify(t, conn, map[string]interface{}{
"tls_v1": true,
"snappy": true,
}, frameTypeResponse)
r := struct {
TLSv1 bool `json:"tls_v1"`
Snappy bool `json:"snappy"`
}{}
err = json.Unmarshal(data, &r)
test.Nil(t, err)
test.Equal(t, true, r.TLSv1)
test.Equal(t, true, r.Snappy)
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
}
tlsConn := tls.Client(conn, tlsConfig)
err = tlsConn.Handshake()
test.Nil(t, err)
resp, _ := nsq.ReadResponse(tlsConn)
frameType, data, _ := nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
compressConn := snappystream.NewReader(tlsConn, snappystream.SkipVerifyChecksum)
resp, _ = nsq.ReadResponse(compressConn)
frameType, data, _ = nsq.UnpackResponse(resp)
t.Logf("frameType: %d, data: %s", frameType, data)
test.Equal(t, frameTypeResponse, frameType)
test.Equal(t, []byte("OK"), data)
}
func TestClientMsgTimeout(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.QueueScanRefreshInterval = 100 * time.Millisecond
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
topicName := "test_cmsg_timeout" + strconv.Itoa(int(time.Now().Unix()))
topic := nsqd.GetTopic(topicName)
msg := NewMessage(<-nsqd.idChan, make([]byte, 100))
topic.PutMessage(msg)
// without this the race detector thinks there's a write
// to msg.Attempts that races with the read in the protocol's messagePump...
// it does not reflect a realistically possible condition
topic.PutMessage(NewMessage(<-nsqd.idChan, make([]byte, 100)))
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, map[string]interface{}{
"msg_timeout": 1000,
}, frameTypeResponse)
sub(t, conn, topicName, "ch")
_, err = nsq.Ready(1).WriteTo(conn)
test.Nil(t, err)
resp, _ := nsq.ReadResponse(conn)
_, data, _ := nsq.UnpackResponse(resp)
msgOut, err := decodeMessage(data)
test.Equal(t, msg.ID, msgOut.ID)
test.Equal(t, msg.Body, msgOut.Body)
_, err = nsq.Ready(0).WriteTo(conn)
test.Nil(t, err)
time.Sleep(1100 * time.Millisecond)
_, err = nsq.Finish(nsq.MessageID(msgOut.ID)).WriteTo(conn)
test.Nil(t, err)
resp, _ = nsq.ReadResponse(conn)
frameType, data, _ := nsq.UnpackResponse(resp)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, fmt.Sprintf("E_FIN_FAILED FIN %s failed ID not in flight", msgOut.ID),
string(data))
}
func TestBadFin(t *testing.T) {
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
identify(t, conn, map[string]interface{}{}, frameTypeResponse)
sub(t, conn, "test_fin", "ch")
fin := nsq.Finish(nsq.MessageID{})
fin.Params[0] = []byte("")
_, err = fin.WriteTo(conn)
test.Nil(t, err)
resp, _ := nsq.ReadResponse(conn)
frameType, data, _ := nsq.UnpackResponse(resp)
test.Equal(t, frameTypeError, frameType)
test.Equal(t, "E_INVALID Invalid Message ID", string(data))
}
func TestClientAuth(t *testing.T) {
authResponse := `{"ttl":1, "authorizations":[]}`
authSecret := "testsecret"
authError := "E_UNAUTHORIZED AUTH No authorizations found"
authSuccess := ""
runAuthTest(t, authResponse, authSecret, authError, authSuccess)
// now one that will succeed
authResponse = `{"ttl":10, "authorizations":
[{"topic":"test", "channels":[".*"], "permissions":["subscribe","publish"]}]
}`
authError = ""
authSuccess = `{"identity":"","identity_url":"","permission_count":1}`
runAuthTest(t, authResponse, authSecret, authError, authSuccess)
}
func runAuthTest(t *testing.T, authResponse, authSecret, authError, authSuccess string) {
var err error
var expectedAuthIP string
expectedAuthTLS := "false"
authd := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Logf("in test auth handler %s", r.RequestURI)
r.ParseForm()
test.Equal(t, expectedAuthIP, r.Form.Get("remote_ip"))
test.Equal(t, expectedAuthTLS, r.Form.Get("tls"))
test.Equal(t, authSecret, r.Form.Get("secret"))
fmt.Fprint(w, authResponse)
}))
defer authd.Close()
addr, err := url.Parse(authd.URL)
test.Nil(t, err)
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
opts.AuthHTTPAddresses = []string{addr.Host}
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
defer nsqd.Exit()
conn, err := mustConnectNSQD(tcpAddr)
test.Nil(t, err)
defer conn.Close()
expectedAuthIP, _, _ = net.SplitHostPort(conn.LocalAddr().String())
identify(t, conn, map[string]interface{}{
"tls_v1": false,
}, nsq.FrameTypeResponse)
authCmd(t, conn, authSecret, authSuccess)
if authError != "" {
readValidate(t, conn, nsq.FrameTypeError, authError)
} else {
sub(t, conn, "test", "ch")
}
}
func TestIOLoopReturnsClientErrWhenSendFails(t *testing.T) {
fakeConn := test.NewFakeNetConn()
fakeConn.WriteFunc = func(b []byte) (int, error) {
return 0, errors.New("write error")
}
testIOLoopReturnsClientErr(t, fakeConn)
}
func TestIOLoopReturnsClientErrWhenSendSucceeds(t *testing.T) {
fakeConn := test.NewFakeNetConn()
fakeConn.WriteFunc = func(b []byte) (int, error) {
return len(b), nil
}
testIOLoopReturnsClientErr(t, fakeConn)
}
func testIOLoopReturnsClientErr(t *testing.T, fakeConn test.FakeNetConn) {
fakeConn.ReadFunc = func(b []byte) (int, error) {
return copy(b, []byte("INVALID_COMMAND\n")), nil
}
opts := NewOptions()
opts.Logger = test.NewTestLogger(t)
opts.Verbose = true
prot := &protocolV2{ctx: &context{nsqd: New(opts)}}
defer prot.ctx.nsqd.Exit()
err := prot.IOLoop(fakeConn)
test.NotNil(t, err)
test.Equal(t, "E_INVALID invalid command INVALID_COMMAND", err.Error())
test.NotNil(t, err.(*protocol.FatalClientErr))
}
func BenchmarkProtocolV2Exec(b *testing.B) {
b.StopTimer()
opts := NewOptions()
opts.Logger = test.NewTestLogger(b)
nsqd := New(opts)
ctx := &context{nsqd}
p := &protocolV2{ctx}
c := newClientV2(0, nil, ctx)
params := [][]byte{[]byte("NOP")}
b.StartTimer()
for i := 0; i < b.N; i++ {
p.Exec(c, params)
}
}
func benchmarkProtocolV2Pub(b *testing.B, size int) {
var wg sync.WaitGroup
b.StopTimer()
opts := NewOptions()
batchSize := int(opts.MaxBodySize) / (size + 4)
opts.Logger = test.NewTestLogger(b)
opts.MemQueueSize = int64(b.N)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
msg := make([]byte, size)
batch := make([][]byte, batchSize)
for i := range batch {
batch[i] = msg
}
topicName := "bench_v2_pub" + strconv.Itoa(int(time.Now().Unix()))
b.SetBytes(int64(len(msg)))
b.StartTimer()
for j := 0; j < runtime.GOMAXPROCS(0); j++ {
wg.Add(1)
go func() {
conn, err := mustConnectNSQD(tcpAddr)
if err != nil {
panic(err.Error())
}
rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
num := b.N / runtime.GOMAXPROCS(0) / batchSize
for i := 0; i < num; i++ {
cmd, _ := nsq.MultiPublish(topicName, batch)
_, err := cmd.WriteTo(rw)
if err != nil {
panic(err.Error())
}
err = rw.Flush()
if err != nil {
panic(err.Error())
}
resp, err := nsq.ReadResponse(rw)
if err != nil {
panic(err.Error())
}
_, data, _ := nsq.UnpackResponse(resp)
if !bytes.Equal(data, []byte("OK")) {
panic("invalid response")
}
}
wg.Done()
}()
}
wg.Wait()
b.StopTimer()
nsqd.Exit()
}
func BenchmarkProtocolV2Pub256(b *testing.B) { benchmarkProtocolV2Pub(b, 256) }
func BenchmarkProtocolV2Pub512(b *testing.B) { benchmarkProtocolV2Pub(b, 512) }
func BenchmarkProtocolV2Pub1k(b *testing.B) { benchmarkProtocolV2Pub(b, 1024) }
func BenchmarkProtocolV2Pub2k(b *testing.B) { benchmarkProtocolV2Pub(b, 2*1024) }
func BenchmarkProtocolV2Pub4k(b *testing.B) { benchmarkProtocolV2Pub(b, 4*1024) }
func BenchmarkProtocolV2Pub8k(b *testing.B) { benchmarkProtocolV2Pub(b, 8*1024) }
func BenchmarkProtocolV2Pub16k(b *testing.B) { benchmarkProtocolV2Pub(b, 16*1024) }
func BenchmarkProtocolV2Pub32k(b *testing.B) { benchmarkProtocolV2Pub(b, 32*1024) }
func BenchmarkProtocolV2Pub64k(b *testing.B) { benchmarkProtocolV2Pub(b, 64*1024) }
func BenchmarkProtocolV2Pub128k(b *testing.B) { benchmarkProtocolV2Pub(b, 128*1024) }
func BenchmarkProtocolV2Pub256k(b *testing.B) { benchmarkProtocolV2Pub(b, 256*1024) }
func BenchmarkProtocolV2Pub512k(b *testing.B) { benchmarkProtocolV2Pub(b, 512*1024) }
func BenchmarkProtocolV2Pub1m(b *testing.B) { benchmarkProtocolV2Pub(b, 1024*1024) }
func benchmarkProtocolV2Sub(b *testing.B, size int) {
var wg sync.WaitGroup
b.StopTimer()
opts := NewOptions()
opts.Logger = test.NewTestLogger(b)
opts.MemQueueSize = int64(b.N)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
msg := make([]byte, size)
topicName := "bench_v2_sub" + strconv.Itoa(b.N) + strconv.Itoa(int(time.Now().Unix()))
topic := nsqd.GetTopic(topicName)
for i := 0; i < b.N; i++ {
msg := NewMessage(<-nsqd.idChan, msg)
topic.PutMessage(msg)
}
topic.GetChannel("ch")
b.SetBytes(int64(len(msg)))
goChan := make(chan int)
rdyChan := make(chan int)
workers := runtime.GOMAXPROCS(0)
for j := 0; j < workers; j++ {
wg.Add(1)
go func() {
subWorker(b.N, workers, tcpAddr, topicName, rdyChan, goChan)
wg.Done()
}()
<-rdyChan
}
b.StartTimer()
close(goChan)
wg.Wait()
b.StopTimer()
nsqd.Exit()
}
func subWorker(n int, workers int, tcpAddr *net.TCPAddr, topicName string, rdyChan chan int, goChan chan int) {
conn, err := mustConnectNSQD(tcpAddr)
if err != nil {
panic(err.Error())
}
rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriterSize(conn, 65536))
identify(nil, conn, nil, frameTypeResponse)
sub(nil, conn, topicName, "ch")
rdyCount := int(math.Min(math.Max(float64(n/workers), 1), 2500))
rdyChan <- 1
<-goChan
nsq.Ready(rdyCount).WriteTo(rw)
rw.Flush()
num := n / workers
for i := 0; i < num; i++ {
resp, err := nsq.ReadResponse(rw)
if err != nil {
panic(err.Error())
}
frameType, data, err := nsq.UnpackResponse(resp)
if err != nil {
panic(err.Error())
}
if frameType != frameTypeMessage {
panic("got something else")
}
msg, err := decodeMessage(data)
if err != nil {
panic(err.Error())
}
nsq.Finish(nsq.MessageID(msg.ID)).WriteTo(rw)
if (i+1)%rdyCount == 0 || i+1 == num {
if i+1 == num {
nsq.Ready(0).WriteTo(conn)
}
rw.Flush()
}
}
}
func BenchmarkProtocolV2Sub256(b *testing.B) { benchmarkProtocolV2Sub(b, 256) }
func BenchmarkProtocolV2Sub512(b *testing.B) { benchmarkProtocolV2Sub(b, 512) }
func BenchmarkProtocolV2Sub1k(b *testing.B) { benchmarkProtocolV2Sub(b, 1024) }
func BenchmarkProtocolV2Sub2k(b *testing.B) { benchmarkProtocolV2Sub(b, 2*1024) }
func BenchmarkProtocolV2Sub4k(b *testing.B) { benchmarkProtocolV2Sub(b, 4*1024) }
func BenchmarkProtocolV2Sub8k(b *testing.B) { benchmarkProtocolV2Sub(b, 8*1024) }
func BenchmarkProtocolV2Sub16k(b *testing.B) { benchmarkProtocolV2Sub(b, 16*1024) }
func BenchmarkProtocolV2Sub32k(b *testing.B) { benchmarkProtocolV2Sub(b, 32*1024) }
func BenchmarkProtocolV2Sub64k(b *testing.B) { benchmarkProtocolV2Sub(b, 64*1024) }
func BenchmarkProtocolV2Sub128k(b *testing.B) { benchmarkProtocolV2Sub(b, 128*1024) }
func BenchmarkProtocolV2Sub256k(b *testing.B) { benchmarkProtocolV2Sub(b, 256*1024) }
func BenchmarkProtocolV2Sub512k(b *testing.B) { benchmarkProtocolV2Sub(b, 512*1024) }
func BenchmarkProtocolV2Sub1m(b *testing.B) { benchmarkProtocolV2Sub(b, 1024*1024) }
func benchmarkProtocolV2MultiSub(b *testing.B, num int) {
var wg sync.WaitGroup
b.StopTimer()
opts := NewOptions()
opts.Logger = test.NewTestLogger(b)
opts.MemQueueSize = int64(b.N)
tcpAddr, _, nsqd := mustStartNSQD(opts)
defer os.RemoveAll(opts.DataPath)
msg := make([]byte, 256)
b.SetBytes(int64(len(msg) * num))
goChan := make(chan int)
rdyChan := make(chan int)
workers := runtime.GOMAXPROCS(0)
for i := 0; i < num; i++ {
topicName := "bench_v2" + strconv.Itoa(b.N) + "_" + strconv.Itoa(i) + "_" + strconv.Itoa(int(time.Now().Unix()))
topic := nsqd.GetTopic(topicName)
for i := 0; i < b.N; i++ {
msg := NewMessage(<-nsqd.idChan, msg)
topic.PutMessage(msg)
}
topic.GetChannel("ch")
for j := 0; j < workers; j++ {
wg.Add(1)
go func() {
subWorker(b.N, workers, tcpAddr, topicName, rdyChan, goChan)
wg.Done()
}()
<-rdyChan
}
}
b.StartTimer()
close(goChan)
wg.Wait()
b.StopTimer()
nsqd.Exit()
}
func BenchmarkProtocolV2MultiSub1(b *testing.B) { benchmarkProtocolV2MultiSub(b, 1) }
func BenchmarkProtocolV2MultiSub2(b *testing.B) { benchmarkProtocolV2MultiSub(b, 2) }
func BenchmarkProtocolV2MultiSub4(b *testing.B) { benchmarkProtocolV2MultiSub(b, 4) }
func BenchmarkProtocolV2MultiSub8(b *testing.B) { benchmarkProtocolV2MultiSub(b, 8) }
func BenchmarkProtocolV2MultiSub16(b *testing.B) { benchmarkProtocolV2MultiSub(b, 16) }