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

Code before:
		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 {

Review comment (suggestion):
if the error occurs somewhere else, this go-routine running `messageRouter()` wont be cleaned up for the connection that is now closed.

I'm not re-thinking where `transactionCleanup()` should go.  perhaps the original setup where it was deferred inside `messageRouter()` was fine, we only need to make sure that `messageRouter()` properly exits and cleans up whenever `close()` is called?  This can be done by re-setting `exitChan` whenever `ConnectToNSQ()` is called and `close(exitChan)` inside of `close()`.


Code after:
	var err error
	defer this.transactionCleanup()
	atomic.StoreInt32(&this.transactionStat, 1)
	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