Repository: nsqio/nsq
File: nsq/writer.go (Go)

Code before:
package nsq

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// not sure identify always return peerinfo
type Writer struct {
	net.Conn
	sync.Mutex
	Info            PeerInfo
	transactionChan chan *writerTransaction
	dataChan        chan []byte
	WriteTimeout    time.Duration
	transactions    []*writerTransaction
	state           int32
	exitChan        chan int
}

type writerTransaction struct {
	doneChan  chan int
	cmd       *Command
	frameType int32
	data      []byte
	err       error
}

func NewWriter() *Writer {
	this := &Writer{
		transactionChan: make(chan *writerTransaction),
		exitChan:        make(chan int),
		dataChan:        make(chan []byte),
		state:           StateDisconnected,
		WriteTimeout:    time.Second,
	}
	return this
}

func (this *Writer) Stop() {
	this.Close()
	if !atomic.CompareAndSwapInt32(&this.state,
		StateConnected, StateDisconnected) {
		return
	}
	close(this.exitChan)
}

func (this *Writer) Write(cmd *Command) (int32, []byte, error) {
	if atomic.LoadInt32(&this.state) != StateConnected {
		return -1, nil, errors.New("Not connected")
	}
	t := &writerTransaction{
		cmd:       cmd,
		frameType: -1,
		doneChan:  make(chan int),
	}
	this.transactionChan <- t
	<-t.doneChan
	return t.frameType, t.data, t.err
}

func (this *Writer) Identify(data map[string]interface{}) error {
	cmd, err := Identify(data)
	if err != nil {
		return err
	}
	_, resp, err := this.Write(cmd)
	if err != nil {
		return err
	}
	return json.Unmarshal(resp, &this.Info)
}

func (this *Writer) ConnectToNSQ(addr string) error {
	this.Lock()
	defer this.Lock()
	if !atomic.CompareAndSwapInt32(&this.state,
		StateDisconnected, StateConnected) {
		return nil
	}
	var err error
	this.Conn, err = net.DialTimeout("tcp", addr, time.Second*5)
	if err != nil {
		return errors.New("fail to connect nsq")
	}
	this.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
	if _, err := this.Conn.Write(MagicV2); err != nil {
		this.Close()
		return err
	}
	atomic.StoreInt32(&this.state, StateConnected)
	go this.readloop()
	go this.messageRouter()
	return nil
}

func (this *Writer) messageRouter() {
	var err error
	defer this.transactionCleanup()
	for {
		select {
		case t := <-this.transactionChan:
			this.transactions = append(this.transactions, t)
			this.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
			if err = t.cmd.Write(this.Conn); err != nil {
				log.Println("write error", err)
				this.Close()
				return
			}
		case buf := <-this.dataChan:
			frameType, data, err := UnpackResponse(buf)
			if frameType == FrameTypeResponse &&
				bytes.Equal(data, []byte("_heartbeat_")) {
				log.Println("received heartbeat from nsqd")
				if err := this.heartbeat(); err != nil {
					this.Close()
					return
				}
			} else {
				t := this.transactions[0]
				this.transactions = this.transactions[1:]
				t.frameType = frameType
				t.data = data
				t.err = err
				t.doneChan <- 1
			}
			if err != nil {
				log.Println("UnpackResponse failed", err)
				return
			}
		case <-this.exitChan:
			return
		}
	}
}

//send heartbeat
func (this *Writer) heartbeat() error {
	this.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
	if err := Nop().Write(this.Conn); err != nil {
		return err
	}
	return nil
}

// cleanup transactions
func (this *Writer) transactionCleanup() {
	this.Lock()
	defer this.Lock()
	for _, t := range this.transactions {
		t.err = errors.New("Not connected")
		t.doneChan <- 1
	}
	this.transactions = this.transactions[:0]
}
func (this *Writer) readloop() {
	rbuf := bufio.NewReader(this.Conn)
	for {
		resp, err := ReadResponse(rbuf)
		if err != nil {
			this.Close()
			break
		}
		select {
		case this.dataChan <- resp:
		case <-this.exitChan:
			return
		}
	}
}


Review comment (bug):
we should log something here, in the format:

``` go
log.Printf("ERROR: failed to ReadResponse from %s - %s", this.Conn, err.Error())
```


Code after:
//send heartbeat
func (this *Writer) heartbeat() error {
	this.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
	if err := Nop().Write(this.Conn); err != nil {
		return err
	}
	return nil
}

// cleanup transactions
func (this *Writer) transactionCleanup() {
	this.Lock()
	defer this.Unlock()
	for _, t := range this.transactions {
		t.err = errors.New("not connected")
		t.doneChan <- 1
	}
	this.transactions = this.transactions[:0]
	atomic.CompareAndSwapInt32(&this.state, StateConnected, StateDisconnected)
}

func (this *Writer) readLoop() {
	rbuf := bufio.NewReader(this.Conn)
	for {
		resp, err := ReadResponse(rbuf)
		if err != nil {
			log.Printf("[%s] error reading response %s", this.RemoteAddr(), err.Error())
			if !strings.Contains(err.Error(), "use of closed network connection") {
				this.Close()
			}
			break
		}
		select {
		case this.dataChan <- resp:
		case <-this.exitChan:
			return
		}
	}
}