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

Code before:
			if cInFlight == 0 {
				log.Printf("[%s] closing FinishedMessages channel", c.ServerAddress)
				close(c.FinishedMessages)
			} else {
				log.Printf("[%s] delaying close of FinishedMesages channel; %d outstanding messages", c.ServerAddress, c.messagesInFlight)
			}
			log.Printf("[%s] stopped read loop ", c.ServerAddress)
			break
		}

		resp, err := c.consumer.ReadResponse()
		if err != nil {
			log.Printf("[%s] error on response %s", c.ServerAddress, err.Error())
			c.stopFlag = true
			continue
		}

		frameType, data, err := c.consumer.UnpackResponse(resp)
		if err != nil {
			log.Printf("[%s] error (%s) unpacking response %s %s", c.ServerAddress, err.Error(), frameType, data)
			c.stopFlag = true
			continue
		}

		switch frameType {
		case FrameTypeMessage:
			c.messagesReceived += 1
			q.MessagesReceived += 1
			msg := data.(*Message)
			log.Printf("[%s] FrameTypeMessage: %s - %s", c.ServerAddress, msg.Id, msg.Body)
			q.Lock()
			c.messagesInFlight += 1
			q.messagesInFlight += 1
			q.Unlock()
			q.IncomingMessages <- &IncomingMessage{msg, c.FinishedMessages}
		case FrameTypeCloseWait:
			// server is ready for us to close (it ack'd our StartClose)
			// we can assume we will not receive any more messages over this channel
			// (but we can still write back responses)
			log.Printf("[%s] received ACK from server. now in CloseWait %d", c.ServerAddress, frameType)
			c.stopFlag = true
		default:
			log.Printf("[%s] unknown message type %d", c.ServerAddress, frameType)
		}

	}
}

func ConnectionFinishLoop(q *Reader, c *nsqConn) {
	for {
		msg, ok := <-c.FinishedMessages

Review comment (suggestion):
in this case, we actually need to go through the loop one more time with the stopFlag set to true so that we will hit the block that closes the FinishedMessages channel and triggers the rest of the exit logic. Not sure that this turns out to be the best approach, I was just trying to make a minimal change


Code after:
		q.stopFlag = true
	}
}

func ConnectionReadLoop(q *Reader, c *nsqConn) {
	// prime our ready state
	c.consumer.WriteCommand(c.consumer.Ready(q.getBufferSize()))
	for {
		if c.stopFlag || q.stopFlag {
			// start the connection close
			q.RLock()
			cInFlight := c.messagesInFlight
			q.RUnlock()
			if cInFlight == 0 {
				log.Printf("[%s] closing FinishedMessages channel", c.ServerAddress)
				close(c.FinishedMessages)
			} else {
				log.Printf("[%s] delaying close of FinishedMesages channel; %d outstanding messages", c.ServerAddress, c.messagesInFlight)
			}
			log.Printf("[%s] stopped read loop ", c.ServerAddress)
			break
		}

		resp, err := c.consumer.ReadResponse()
		if err != nil {
			handleReadError(q, c, fmt.Sprintf("[%s] error on response %s", c.ServerAddress, err.Error()))
			continue
		}

		frameType, data, err := c.consumer.UnpackResponse(resp)
		if err != nil {
			handleReadError(q, c, fmt.Sprintf("[%s] error (%s) unpacking response %s %s", c.ServerAddress, err.Error(), frameType, data))
			continue
		}

		switch frameType {
		case FrameTypeMessage:
			c.messagesReceived += 1
			q.MessagesReceived += 1
			msg := data.(*Message)
			log.Printf("[%s] FrameTypeMessage: %s - %s", c.ServerAddress, msg.Id, msg.Body)
			q.Lock()
			c.messagesInFlight += 1
			q.messagesInFlight += 1
			q.Unlock()
			q.IncomingMessages <- &IncomingMessage{msg, c.FinishedMessages}
		case FrameTypeCloseWait:
			// server is ready for us to close (it ack'd our StartClose)
			// we can assume we will not receive any more messages over this channel
			// (but we can still write back responses)
			log.Printf("[%s] received ACK from server. now in CloseWait %d", c.ServerAddress, frameType)