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

Code before:
package nsq

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// define internal write message struct
type WMessage struct {
	Cmd  *Command
	Stat chan error
}

// may be should add more info in writer
type Writer struct {
	net.Conn
	MessageChan       chan WMessage
	transactionChan   chan error
	ShortIdentifier   string
	LongIdentifier    string
	HeartBeatInterval int
	WriteTimeout      time.Duration
	VerboseLogging    bool
	ErrorStat         chan error
}

func NewWriter() *Writer {
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("ERROR: unable to get hostname %s", err.Error())
	}
	this := &Writer{
		MessageChan:       make(chan WMessage),
		transactionChan:   make(chan error),
		ErrorStat:         make(chan error),
		ShortIdentifier:   strings.Split(hostname, ".")[0],
		LongIdentifier:    hostname,
		WriteTimeout:      time.Second,
		HeartBeatInterval: 30000,
		VerboseLogging:    false,
	}
	return this
}

// stop writerloop
func (this *Writer) Stop() {
	close(this.MessageChan)
}

// body should be []byte or [][]byte
func (this *Writer) Write(topic string, body interface{}) (err error) {
	defer func() {
		if x := recover(); x != nil {
			err = fmt.Errorf("Write Closed: %v", x)
		}
	}()
	var cmd *Command
	switch body.(type) {
	case []byte:
		cmd = Publish(topic, body.([]byte))
	case [][]byte:
		cmd, err = MultiPublish(topic, body.([][]byte))
	case error:
		cmd = Nop()
	default:
		return errors.New("Body not acceptable")
	}
	msg := WMessage{
		Cmd:  cmd,
		Stat: make(chan error),
	}
	this.MessageChan <- msg
	err = <-msg.Stat
	return
}

// always retry to connect nsqd
func (this *Writer) ConnectToNSQ(addr string) error {
	var err error
	this.Conn, err = net.DialTimeout("tcp", addr, time.Second*5)
	if err != nil {
		return errors.New("fail to connect nsq")
	}
	this.Conn.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
	this.Conn.Write(MagicV2)
	if err := this.identify(); err != nil {
		this.Conn.Close()
		return err
	}
	go this.writeloop()
	go this.readloop()
	go this.retry_connect(addr)
	return nil
}

func (this *Writer) identify() error {
	ci := make(map[string]interface{})
	ci["short_id"] = this.ShortIdentifier
	ci["long_id"] = this.LongIdentifier
	ci["heartbeat_interval"] = this.HeartBeatInterval
	cmd, err := Identify(ci)
	if err = cmd.Write(this.Conn); err != nil {
		log.Println("write identify failed", err)
		return err
	}
	rbuf := bufio.NewReader(this.Conn)
	resp, err := ReadResponse(rbuf)
	if err != nil {
		return err
	}
	_, data, err := UnpackResponse(resp)
	if err != nil {
		return errors.New("unpack failed")
	}
	if !bytes.Equal(data, []byte("OK")) {
		return errors.New("fail to identify")
	}
	return nil
}

func (this *Writer) retry_connect(addr string) {
	var err error
	for {
		if err = <-this.ErrorStat; err != nil {
			log.Println(err)
			if this.Conn != nil {
				this.Conn.Close()
			}
			if err.Error() == "ExitLoop" {
				break
			}
		}
		exitchan := make(chan int)
		go this.handle_close(exitchan)
		log.Println("Try to reconnect nsq", addr)
		this.Conn, err = net.DialTimeout("tcp", addr, time.Second*5)
		if err != nil {
			go func() {
				this.ErrorStat <- err
			}()
			time.Sleep(time.Second)
			close(exitchan)
			continue
		}
		this.Conn.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
		this.Conn.Write(MagicV2)
		if err := this.identify(); err != nil {
			go func() {
				this.ErrorStat <- err
			}()
			time.Sleep(time.Second)
			close(exitchan)
			continue
		}
		close(exitchan)
		go this.writeloop()
		go this.readloop()
	}
	log.Println("Writer was Stopped")
}

// handle messge when no connection to nsq
func (this *Writer) handle_close(exitchan chan int) {
	for {
		select {
		case <-exitchan:
			return
		case msg, ok := <-this.MessageChan:
			msg.Stat <- errors.New("no nsq connection")
			if !ok {
				return
			}
		}
	}
}

// write loop until tcp error
func (this *Writer) writeloop() {
	var err error
	log.Println("start writeloop")
	for {
		msg, ok := <-this.MessageChan
		if !ok {
			err = errors.New("ExitLoop")
			break
		}
		this.Conn.SetWriteDeadline(time.Now().Add(this.WriteTimeout))
		if err = msg.Cmd.Write(this.Conn); err != nil {
			log.Println("write error", err)
			msg.Stat <- err
			break
		}
		if string(msg.Cmd.Name) == "NOP" {
			msg.Stat <- nil
		} else {
			err = <-this.transactionChan
			msg.Stat <- err
			if err != nil && err.Error() == "ReadFailed" {
				break
			}
		}
	}
	this.ErrorStat <- err
	log.Println("exit writeloop")
}

func (this *Writer) readloop() {
	rbuf := bufio.NewReader(this.Conn)
	log.Println("start readloop")
	for {
		resp, err := ReadResponse(rbuf)
		if err != nil {
			log.Println("failed to read response", err)
			if err == io.EOF ||
				strings.Contains(err.Error(),
					"use of closed network connection") {
				msg := WMessage{
					Cmd:  Nop(),
					Stat: make(chan error),
				}
				go func() {
					defer func() { recover() }()
					this.MessageChan <- msg
					<-msg.Stat
				}()
			} else {
				this.transactionChan <- errors.New("ReadFailed")
			}
			break
		}
		frametype, data, err := UnpackResponse(resp)
		if err != nil {
			log.Println("unpack failed", err)
			this.transactionChan <- err
			continue
		}
		if bytes.Equal(data, []byte("OK")) {
			this.transactionChan <- nil
			continue
		}
		if bytes.Equal(data, []byte("_heartbeat_")) &&
			frametype == FrameTypeResponse {
			msg := WMessage{
				Cmd:  Nop(),
				Stat: make(chan error),
			}
			go func() {
				defer func() { recover() }()
				this.MessageChan <- msg
				<-msg.Stat
			}()
			continue
		}
		log.Println("response not ok", string(data))
		this.transactionChan <- errors.New(string(data))
	}
	log.Println("exit readloop")
}


Review comment (suggestion):
actually, I take this back, this should wait for a response on `<-msg.Stat` (one of the other goroutines should have the responsibility of making sure all pending transactions get answered)


Code after:
package nsq

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Writer struct {
	net.Conn
	sync.Mutex
	transactionChan   chan *writerTransaction
	dataChan          chan []byte
	WriteTimeout      time.Duration
	transactions      []*writerTransaction
	state             int32
	stopFlag          int32
	HeartbeatInterval int
	ShortIdentifier   string
	LongIdentifier    string
	exitChan          chan int
}

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

func NewWriter(heartbeatInterval int) *Writer {
	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("ERROR: unable to get hostname %s", err.Error())
	}
	this := &Writer{
		transactionChan:   make(chan *writerTransaction),
		exitChan:          make(chan int),
		dataChan:          make(chan []byte),
		state:             StateDisconnected,
		stopFlag:          0,
		WriteTimeout:      time.Second,
		HeartbeatInterval: heartbeatInterval,
		ShortIdentifier:   strings.Split(hostname, ".")[0],
		LongIdentifier:    hostname,
	}
	return this
}

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

func (this *Writer) Publish(topic string, body []byte) (int32, []byte, error) {
	cmd := Publish(topic, body)
	return this.sendCommand(cmd)
}

func (this *Writer) MultiPublish(topic string, body [][]byte) (int32, []byte, error) {
	cmd, err := MultiPublish(topic, body)
	if err != nil {
		return -1, nil, err
	}
	return this.sendCommand(cmd)
}

func (this *Writer) sendCommand(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) ConnectToNSQ(addr string) error {
	if atomic.LoadInt32(&this.stopFlag) == 1 {
		return errors.New("writer stopped")
	}
	this.Lock()
	defer this.Unlock()
	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("failed 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)
	ci := make(map[string]interface{})
	ci["short_id"] = this.ShortIdentifier
	ci["long_id"] = this.LongIdentifier
	ci["heartbeat_interval"] = this.HeartbeatInterval
	ci["feature_negotiation"] = true
	cmd, err := Identify(ci)
	if err != nil {
		this.Close()
		return fmt.Errorf("[%s] failed to create identify command - %s", this.RemoteAddr(), err.Error())
	}
	go this.readLoop()
	go this.messageRouter()
	frameType, data, err := this.sendCommand(cmd)
	if frameType == FrameTypeError {
		if err == nil {
			this.Close()
			err = fmt.Errorf("[%s] error failed to identify - %s", this.RemoteAddr(), string(data))
		}
	}
	return err
}

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.Printf("[%s] error writing %s", this.RemoteAddr(), err.Error())
				this.Close()
				return
			}
		case buf := <-this.dataChan:
			frameType, data, err := UnpackResponse(buf)
			if err != nil {
				log.Printf("[%s] error (%s) unpacking response %d %s", this.RemoteAddr(), err.Error(), frameType, data)
				return
			}
			if frameType == FrameTypeResponse &&
				bytes.Equal(data, []byte("_heartbeat_")) {
				log.Printf("[%s] received heartbeat", this.RemoteAddr())
				if err := this.heartbeat(); err != nil {
					log.Printf("[%s] error sending heartbeat - %s", this.RemoteAddr(), err.Error())
					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
			}
		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.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
		}
	}
}