Merge pull request nsqio#121 from mreiferson/buffered_producer_chan_121

producer: buffer transaction chan

command_test.go

@@ -2,16 +2,11 @@ package nsq
 
 import (
 	"bytes"
-	"io/ioutil"
-	"log"
-	"os"
 	"testing"
 )
 
 func BenchmarkCommand(b *testing.B) {
 	b.StopTimer()
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
 	data := make([]byte, 2048)
 	cmd := Publish("test", data)
 	var buf bytes.Buffer

mock_test.go

@@ -10,7 +10,6 @@ import (
 	"io/ioutil"
 	"log"
 	"net"
-	"os"
 	"strconv"
 	"testing"
 	"time"
@@ -182,9 +181,6 @@ func (h *testHandler) HandleMessage(message *Message) error {
 }
 
 func TestConsumerBackoff(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	logger := log.New(ioutil.Discard, "", log.LstdFlags)
 
 	var mgood bytes.Buffer

producer.go

@@ -9,6 +9,14 @@ import (
 	"time"
 )
 
+type producerConn interface {
+	String() string
+	SetLogger(logger, LogLevel, string)
+	Connect() (*IdentifyResponse, error)
+	Close() error
+	WriteCommand(*Command) error
+}
+
 // Producer is a high-level type to publish to NSQ.
 //
 // A Producer instance is 1:1 with a destination `nsqd`
@@ -17,7 +25,7 @@ import (
 type Producer struct {
 	id     int64
 	addr   string
-	conn   *Conn
+	conn   producerConn
 	config Config
 
 	logger   logger

producer_test.go

@@ -1,11 +1,12 @@
 package nsq
 
 import (
+	"bytes"
 	"errors"
-	"io/ioutil"
-	"log"
-	"os"
+	"runtime"
 	"strconv"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -35,9 +36,6 @@ func (h *ConsumerHandler) HandleMessage(message *Message) error {
 }
 
 func TestProducerConnection(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	config := NewConfig()
 	w, _ := NewProducer("127.0.0.1:4150", config)
 	w.SetLogger(nullLogger, LogLevelInfo)
@@ -56,9 +54,6 @@ func TestProducerConnection(t *testing.T) {
 }
 
 func TestProducerPublish(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	topicName := "publish" + strconv.Itoa(int(time.Now().Unix()))
 	msgCount := 10
 
@@ -83,9 +78,6 @@ func TestProducerPublish(t *testing.T) {
 }
 
 func TestProducerMultiPublish(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	topicName := "multi_publish" + strconv.Itoa(int(time.Now().Unix()))
 	msgCount := 10
 
@@ -113,9 +105,6 @@ func TestProducerMultiPublish(t *testing.T) {
 }
 
 func TestProducerPublishAsync(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	topicName := "async_publish" + strconv.Itoa(int(time.Now().Unix()))
 	msgCount := 10
 
@@ -151,9 +140,6 @@ func TestProducerPublishAsync(t *testing.T) {
 }
 
 func TestProducerMultiPublishAsync(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	topicName := "multi_publish" + strconv.Itoa(int(time.Now().Unix()))
 	msgCount := 10
 
@@ -193,9 +179,6 @@ func TestProducerMultiPublishAsync(t *testing.T) {
 }
 
 func TestProducerHeartbeat(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
-	defer log.SetOutput(os.Stdout)
-
 	topicName := "heartbeat" + strconv.Itoa(int(time.Now().Unix()))
 
 	config := NewConfig()
@@ -269,3 +252,85 @@ func readMessages(topicName string, t *testing.T, msgCount int) {
 		t.Fatal("failed message not done")
 	}
 }
+
+type mockProducerConn struct {
+	delegate ConnDelegate
+	closeCh  chan struct{}
+	pubCh    chan struct{}
+}
+
+func newMockProducerConn(delegate ConnDelegate) producerConn {
+	m := &mockProducerConn{
+		delegate: delegate,
+		closeCh:  make(chan struct{}),
+		pubCh:    make(chan struct{}, 4),
+	}
+	go m.router()
+	return m
+}
+
+func (m *mockProducerConn) String() string {
+	return "127.0.0.1:0"
+}
+
+func (m *mockProducerConn) SetLogger(logger logger, level LogLevel, prefix string) {}
+
+func (m *mockProducerConn) Connect() (*IdentifyResponse, error) {
+	return &IdentifyResponse{}, nil
+}
+
+func (m *mockProducerConn) Close() error {
+	close(m.closeCh)
+	return nil
+}
+
+func (m *mockProducerConn) WriteCommand(cmd *Command) error {
+	if bytes.Equal(cmd.Name, []byte("PUB")) {
+		m.pubCh <- struct{}{}
+	}
+	return nil
+}
+
+func (m *mockProducerConn) router() {
+	for {
+		select {
+		case <-m.closeCh:
+			goto exit
+		case <-m.pubCh:
+			m.delegate.OnResponse(nil, framedResponse(FrameTypeResponse, []byte("OK")))
+		}
+	}
+exit:
+}
+
+func BenchmarkProducer(b *testing.B) {
+	b.StopTimer()
+	body := make([]byte, 512)
+
+	config := NewConfig()
+	p, _ := NewProducer("127.0.0.1:0", config)
+
+	p.conn = newMockProducerConn(&producerConnDelegate{p})
+	atomic.StoreInt32(&p.state, StateConnected)
+	p.closeChan = make(chan int)
+	go p.router()
+
+	startCh := make(chan struct{})
+	var wg sync.WaitGroup
+	parallel := runtime.GOMAXPROCS(0)
+
+	for j := 0; j < parallel; j++ {
+		wg.Add(1)
+		go func() {
+			<-startCh
+			for i := 0; i < b.N/parallel; i++ {
+				p.Publish("test", body)
+			}
+			wg.Done()
+		}()
+	}
+
+	b.StartTimer()
+	close(startCh)
+	wg.Wait()
+}