Repository: nsqio/nsq
File: nsqd/protocol_v2.go (Go)

Code before:
package main

import (
	"../nsq"
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

const maxTimeout = time.Hour
const maxReadyCount = 2500

type ProtocolV2 struct {
	nsq.Protocol
}

func init() {
	// BigEndian client byte sequence "  V2"
	var magicInt int32
	buf := bytes.NewBuffer([]byte(nsq.MagicV2))
	binary.Read(buf, binary.BigEndian, &magicInt)
	protocols[magicInt] = &ProtocolV2{}
}

func (p *ProtocolV2) IOLoop(conn net.Conn) error {
	var err error
	var line string

	client := NewClientV2(conn)
	atomic.StoreInt32(&client.State, nsq.StateInit)

	err = nil
	reader := bufio.NewReader(client)
	for {
		client.SetReadDeadline(time.Now().Add(nsqd.clientTimeout))
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}

		line = strings.TrimSpace(line)
		params := strings.Split(line, " ")

		if *verbose {
			log.Printf("PROTOCOL(V2): [%s] %#v", client.RemoteAddr().String(), params)
		}

		response, err := p.Exec(client, params)
		if err != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeError, []byte(err.Error()))
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
			continue
		}

		if response != nil {
			clientData, err := nsq.Frame(nsq.FrameTypeResponse, response)
			if err != nil {
				break
			}

			err = p.Write(client, clientData)
			if err != nil {
				break
			}
		}
	}

	log.Printf("PROTOCOL(V2): [%s] exiting ioloop", client.RemoteAddr().String())
	// TODO: gracefully send clients the close signal
	conn.Close()
	close(client.ExitChan)

	return err
}

func (p *ProtocolV2) Write(client *ClientV2, data []byte) error {
	attempts := 0
	for {
		client.Lock()
		client.SetWriteDeadline(time.Now().Add(time.Second))
		_, err := nsq.SendResponse(client, data)
		client.Unlock()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				attempts++
				if attempts == 3 {
					return errors.New("E_WRITE_TIMEOUT")
				}
				continue
			}
			return err
		}
		break
	}
	return nil
}

func (p *ProtocolV2) Exec(client *ClientV2, params []string) ([]byte, error) {
	switch params[0] {
	case "SUB":
		return p.SUB(client, params)
	case "RDY":
		return p.RDY(client, params)
	case "FIN":
		return p.FIN(client, params)
	case "REQ":
		return p.REQ(client, params)
	case "CLS":
		return p.CLS(client, params)
	case "NOP":
		return p.NOP(client, params)
	}
	return nil, nsq.ClientErrInvalid
}

func (p *ProtocolV2) sendHeartbeat(client *ClientV2) error {
	clientData, err := nsq.Frame(nsq.FrameTypeResponse, []byte("_heartbeat_"))
	if err != nil {
		return err
	}

	err = p.Write(client, clientData)
	if err != nil {
		return err
	}

	return nil
}

func (p *ProtocolV2) messagePump(client *ClientV2) {
	var err error

	heartbeat := time.NewTicker(nsqd.clientTimeout / 2)

	// ReadyStateChan has a buffer of 1 to guarantee that in the event
	// there is a race before we enter either of the select loops where 
	// ReadyStateChan is read from that the update is not lost
	for {
		if !client.IsReadyForMessages() {
			// wait for a change in state
			select {
			case <-client.ReadyStateChan:
			case <-heartbeat.C:
				err := p.sendHeartbeat(client)
				if err != nil {
					log.Printf("PROTOCOL(V2): error sending heartbeat - %s", err.Error())
				}
			case <-client.ExitChan:
				goto exit
			}
		} else {
			select {
			case <-client.ReadyStateChan:
			case <-heartbeat.C:
				err := p.sendHeartbeat(client)
				if err != nil {
					log.Printf("PROTOCOL(V2): error sending heartbeat - %s", err.Error())
				}
			case msg, ok := <-client.Channel.clientMsgChan:
				if !ok {
					goto exit
				}

				if *verbose {
					log.Printf("PROTOCOL(V2): writing msg(%s) to client(%s) - %s",
						msg.Id, client, msg.Body)
				}

				data, err := msg.Encode()
				if err != nil {
					goto exit
				}

				clientData, err := nsq.Frame(nsq.FrameTypeMessage, data)
				if err != nil {
					goto exit
				}

				client.Channel.StartInFlightTimeout(msg, client)
				client.SendingMessage()

				err = p.Write(client, clientData)
				if err != nil {
					goto exit
				}
			case <-client.ExitChan:
				goto exit
			}
		}
	}

exit:
	log.Printf("PROTOCOL(V2): [%s] exiting messagePump", client.RemoteAddr().String())
	heartbeat.Stop()
	client.Channel.RemoveClient(client)
	if err != nil {
		log.Printf("PROTOCOL(V2): messagePump error - %s", err.Error())
	}
}

func (p *ProtocolV2) SUB(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateInit {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	topicName := params[1]
	if !nsq.IsValidTopicName(topicName) {
		return nil, nsq.ClientErrBadTopic
	}

	channelName := params[2]
	if !nsq.IsValidChannelName(channelName) {
		return nil, nsq.ClientErrBadChannel
	}

	if len(params) == 5 {
		client.ShortIdentifier = params[3]
		client.LongIdentifier = params[4]
	}

	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel(channelName)
	channel.AddClient(client)

	client.Channel = channel
	atomic.StoreInt32(&client.State, nsq.StateSubscribed)

	go p.messagePump(client)

	return nil, nil
}

func (p *ProtocolV2) RDY(client *ClientV2, params []string) ([]byte, error) {
	var err error

	state := atomic.LoadInt32(&client.State)

	if state == nsq.StateClosing {
		// just ignore ready changes on a closing channel
		log.Printf("PROTOCOL(V2): ignoring RDY after CLS in state ClientStateV2Closing")
		return nil, nil
	}

	if state != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	count := 1
	if len(params) > 1 {
		count, err = strconv.Atoi(params[1])
		if err != nil {
			return nil, nsq.ClientErrInvalid
		}
	}

	if count > maxReadyCount {
		log.Printf("ERROR: client(%s) sent ready count %d > %d", client, count, maxReadyCount)
		count = maxReadyCount
	}

	client.SetReadyCount(int64(count))

	return nil, nil
}

func (p *ProtocolV2) FIN(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 2 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	err := client.Channel.FinishMessage(client, []byte(idStr))
	if err != nil {
		return nil, nsq.ClientErrFinishFailed
	}

	client.FinishedMessage()

	return nil, nil
}

func (p *ProtocolV2) REQ(client *ClientV2, params []string) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	timeoutMs, err := strconv.Atoi(params[2])
	if err != nil {
		return nil, nsq.ClientErrInvalid
	}
	timeoutDuration := time.Duration(timeoutMs) * time.Millisecond

	if timeoutDuration < 0 || timeoutDuration > maxTimeout {
		return nil, nsq.ClientErrInvalid
	}

	err = client.Channel.RequeueMessage(client, []byte(idStr), timeoutDuration)
	if err != nil {
		return nil, nsq.ClientErrRequeueFailed
	}

	client.RequeuedMessage()

	return nil, nil
}

func (p *ProtocolV2) CLS(client *ClientV2, params []string) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) > 1 {
		return nil, nsq.ClientErrInvalid
	}

	client.StartClose()

	return []byte("CLOSE_WAIT"), nil
}

func (p *ProtocolV2) NOP(client *ClientV2, params []string) ([]byte, error) {
	return nil, nil
}


Review comment (suggestion):
now that we're re-using a buffer when this method is called on the protocol, we need to move where we lock to instead be around the whole method rather than just `nsq.SendResponse`


Code after:
package main

import (
	"../nsq"
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

const maxTimeout = time.Hour
const maxReadyCount = 2500

type ProtocolV2 struct {
	nsq.Protocol
}

func init() {
	// BigEndian client byte sequence "  V2"
	var magicInt int32
	buf := bytes.NewBuffer([]byte(nsq.MagicV2))
	binary.Read(buf, binary.BigEndian, &magicInt)
	protocols[magicInt] = &ProtocolV2{}
}

func (p *ProtocolV2) IOLoop(conn net.Conn) error {
	var err error
	var line []byte

	client := NewClientV2(conn)
	atomic.StoreInt32(&client.State, nsq.StateInit)

	err = nil
	reader := bufio.NewReader(client)
	for {
		client.SetReadDeadline(time.Now().Add(nsqd.clientTimeout))
		// ReadSlice does not allocate new space for the data each request
		// ie. the returned slice is only valid until the next call to it
		line, err = reader.ReadSlice('\n')
		if err != nil {
			break
		}

		// trim the '\n'
		line = line[:len(line)-1]
		// optionally trim the '\r'
		if line[len(line)-1] == '\r' {
			line = line[:len(line)-1]
		}
		params := bytes.Split(line, []byte(" "))

		if *verbose {
			log.Printf("PROTOCOL(V2): [%s] %#v", client.RemoteAddr().String(), params)
		}

		response, err := p.Exec(client, params)
		if err != nil {
			err = p.Send(client, nsq.FrameTypeError, []byte(err.Error()))
			if err != nil {
				break
			}
			continue
		}

		if response != nil {
			err = p.Send(client, nsq.FrameTypeResponse, response)
			if err != nil {
				break
			}
		}
	}

	log.Printf("PROTOCOL(V2): [%s] exiting ioloop", client.RemoteAddr().String())
	// TODO: gracefully send clients the close signal
	conn.Close()
	close(client.ExitChan)

	return err
}

func (p *ProtocolV2) Send(client *ClientV2, frameType int32, data []byte) error {
	client.Lock()
	defer client.Unlock()

	client.frameBuf.Reset()
	err := nsq.Frame(&client.frameBuf, frameType, data)
	if err != nil {
		return err
	}

	attempts := 0
	for {
		client.SetWriteDeadline(time.Now().Add(time.Second))
		_, err := nsq.SendResponse(client, client.frameBuf.Bytes())
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				attempts++
				if attempts == 3 {
					return errors.New("E_WRITE_TIMEOUT")
				}
				continue
			}
			return err
		}
		break
	}

	return nil
}

func (p *ProtocolV2) Exec(client *ClientV2, params [][]byte) ([]byte, error) {
	switch {
	case bytes.Equal(params[0], []byte("SUB")):
		return p.SUB(client, params)
	case bytes.Equal(params[0], []byte("RDY")):
		return p.RDY(client, params)
	case bytes.Equal(params[0], []byte("FIN")):
		return p.FIN(client, params)
	case bytes.Equal(params[0], []byte("REQ")):
		return p.REQ(client, params)
	case bytes.Equal(params[0], []byte("CLS")):
		return p.CLS(client, params)
	case bytes.Equal(params[0], []byte("NOP")):
		return p.NOP(client, params)
	}
	return nil, nsq.ClientErrInvalid
}

func (p *ProtocolV2) sendHeartbeat(client *ClientV2) error {
	err := p.Send(client, nsq.FrameTypeResponse, []byte("_heartbeat_"))
	if err != nil {
		return err
	}

	return nil
}

func (p *ProtocolV2) messagePump(client *ClientV2) {
	var err error
	var buf bytes.Buffer

	heartbeat := time.NewTicker(nsqd.clientTimeout / 2)

	// ReadyStateChan has a buffer of 1 to guarantee that in the event
	// there is a race before we enter either of the select loops where 
	// ReadyStateChan is read from that the update is not lost
	for {
		if !client.IsReadyForMessages() {
			// wait for a change in state
			select {
			case <-client.ReadyStateChan:
			case <-heartbeat.C:
				err = p.sendHeartbeat(client)
				if err != nil {
					log.Printf("PROTOCOL(V2): error sending heartbeat - %s", err.Error())
				}
			case <-client.ExitChan:
				goto exit
			}
		} else {
			select {
			case <-client.ReadyStateChan:
			case <-heartbeat.C:
				err = p.sendHeartbeat(client)
				if err != nil {
					log.Printf("PROTOCOL(V2): error sending heartbeat - %s", err.Error())
				}
			case msg, ok := <-client.Channel.clientMsgChan:
				if !ok {
					goto exit
				}

				if *verbose {
					log.Printf("PROTOCOL(V2): writing msg(%s) to client(%s) - %s",
						msg.Id, client, msg.Body)
				}

				buf.Reset()
				err = msg.Encode(&buf)
				if err != nil {
					goto exit
				}

				client.Channel.StartInFlightTimeout(msg, client)
				client.SendingMessage()

				err = p.Send(client, nsq.FrameTypeMessage, buf.Bytes())
				if err != nil {
					goto exit
				}
			case <-client.ExitChan:
				goto exit
			}
		}
	}

exit:
	log.Printf("PROTOCOL(V2): [%s] exiting messagePump", client.RemoteAddr().String())
	heartbeat.Stop()
	client.Channel.RemoveClient(client)
	if err != nil {
		log.Printf("PROTOCOL(V2): messagePump error - %s", err.Error())
	}
}

func (p *ProtocolV2) SUB(client *ClientV2, params [][]byte) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateInit {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	topicName := string(params[1])
	if !nsq.IsValidTopicName(topicName) {
		return nil, nsq.ClientErrBadTopic
	}

	channelName := string(params[2])
	if !nsq.IsValidChannelName(channelName) {
		return nil, nsq.ClientErrBadChannel
	}

	if len(params) == 5 {
		client.ShortIdentifier = string(params[3])
		client.LongIdentifier = string(params[4])
	}

	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel(channelName)
	channel.AddClient(client)

	client.Channel = channel
	atomic.StoreInt32(&client.State, nsq.StateSubscribed)

	go p.messagePump(client)

	return nil, nil
}

func (p *ProtocolV2) RDY(client *ClientV2, params [][]byte) ([]byte, error) {
	var err error

	state := atomic.LoadInt32(&client.State)

	if state == nsq.StateClosing {
		// just ignore ready changes on a closing channel
		log.Printf("PROTOCOL(V2): ignoring RDY after CLS in state ClientStateV2Closing")
		return nil, nil
	}

	if state != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	count := 1
	if len(params) > 1 {
		count, err = strconv.Atoi(string(params[1]))
		if err != nil {
			return nil, nsq.ClientErrInvalid
		}
	}

	if count > maxReadyCount {
		log.Printf("ERROR: client(%s) sent ready count %d > %d", client, count, maxReadyCount)
		count = maxReadyCount
	}

	client.SetReadyCount(int64(count))

	return nil, nil
}

func (p *ProtocolV2) FIN(client *ClientV2, params [][]byte) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 2 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	err := client.Channel.FinishMessage(client, idStr)
	if err != nil {
		return nil, nsq.ClientErrFinishFailed
	}

	client.FinishedMessage()

	return nil, nil
}

func (p *ProtocolV2) REQ(client *ClientV2, params [][]byte) ([]byte, error) {
	state := atomic.LoadInt32(&client.State)
	if state != nsq.StateSubscribed && state != nsq.StateClosing {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) < 3 {
		return nil, nsq.ClientErrInvalid
	}

	idStr := params[1]
	timeoutMs, err := strconv.Atoi(string(params[2]))
	if err != nil {
		return nil, nsq.ClientErrInvalid
	}
	timeoutDuration := time.Duration(timeoutMs) * time.Millisecond

	if timeoutDuration < 0 || timeoutDuration > maxTimeout {
		return nil, nsq.ClientErrInvalid
	}

	err = client.Channel.RequeueMessage(client, idStr, timeoutDuration)
	if err != nil {
		return nil, nsq.ClientErrRequeueFailed
	}

	client.RequeuedMessage()

	return nil, nil
}

func (p *ProtocolV2) CLS(client *ClientV2, params [][]byte) ([]byte, error) {
	if atomic.LoadInt32(&client.State) != nsq.StateSubscribed {
		return nil, nsq.ClientErrInvalid
	}

	if len(params) > 1 {
		return nil, nsq.ClientErrInvalid
	}

	client.StartClose()

	return []byte("CLOSE_WAIT"), nil
}

func (p *ProtocolV2) NOP(client *ClientV2, params [][]byte) ([]byte, error) {
	return nil, nil
}