// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
var (
b []byte
err error
i, n int
msg message.Message
)
if this.in == nil {
return nil, 0, ErrBufferNotReady
}
// Peek until we get total bytes
for i = 0; ; i++ {
// Peek remlen bytes from the input buffer.
b, err = this.in.ReadWait(total)
if err != nil && err != ErrBufferInsufficientData {
return nil, 0, err
}
// If not enough bytes are returned, then continue until there's enough.
if len(b) >= total {
break
}
}
msg, err = mtype.New()
if err != nil {
return nil, 0, err
}
n, err = msg.Decode(b)
return msg, n, err
}
func writeMessage(conn io.Closer, msg message.Message) error {
buf := make([]byte, msg.Len())
_, err := msg.Encode(buf)
if err != nil {
//glog.Debugf("Write error: %v", err)
return err
}
//glog.Debugf("Writing: %s", msg)
return writeMessageBuffer(conn, buf)
}
// Wait() copies the message into a waiting queue, and waits for the corresponding
// ack message to be received.
func (this *Ackqueue) Wait(msg message.Message, onComplete interface{}) error {
this.mu.Lock()
defer this.mu.Unlock()
switch msg := msg.(type) {
case *message.PublishMessage:
if msg.QoS() == message.QosAtMostOnce {
//return fmt.Errorf("QoS 0 messages don't require ack")
return errWaitMessage
}
this.insert(msg.PacketId(), msg, onComplete)
case *message.SubscribeMessage:
this.insert(msg.PacketId(), msg, onComplete)
case *message.UnsubscribeMessage:
this.insert(msg.PacketId(), msg, onComplete)
case *message.PingreqMessage:
this.ping = ackmsg{
Mtype: message.PINGREQ,
State: message.RESERVED,
OnComplete: onComplete,
}
default:
return errWaitMessage
}
return nil
}
func (this *Ackqueue) insert(pktid uint16, msg message.Message, onComplete interface{}) error {
if this.full() {
this.grow()
}
if _, ok := this.emap[pktid]; !ok {
// message length
ml := msg.Len()
// ackmsg
am := ackmsg{
Mtype: msg.Type(),
State: message.RESERVED,
Pktid: msg.PacketId(),
Msgbuf: make([]byte, ml),
OnComplete: onComplete,
}
if _, err := msg.Encode(am.Msgbuf); err != nil {
return err
}
this.ring[this.tail] = am
this.emap[pktid] = this.tail
this.tail = this.increment(this.tail)
this.count++
} else {
// If packet w/ pktid already exist, then this must be a PUBLISH message
// Other message types should never send with the same packet ID
pm, ok := msg.(*message.PublishMessage)
if !ok {
return fmt.Errorf("ack/insert: duplicate packet ID for %s message", msg.Name())
}
// If this is a publish message, then the DUP flag must be set. This is the
// only scenario in which we will receive duplicate messages.
if pm.Dup() {
return fmt.Errorf("ack/insert: duplicate packet ID for PUBLISH message, but DUP flag is not set")
}
// Since it's a dup, there's really nothing we need to do. Moving on...
}
return nil
}
// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
var (
b []byte
err error
n int
msg message.Message
)
if this.in == nil {
err = ErrBufferNotReady
return nil, 0, err
}
if len(this.intmp) < total {
this.intmp = make([]byte, total)
}
// Read until we get total bytes
l := 0
for l < total {
n, err = this.in.Read(this.intmp[l:])
l += n
glog.Debugf("read %d bytes, total %d", n, l)
if err != nil {
return nil, 0, err
}
}
b = this.intmp[:total]
msg, err = mtype.New()
if err != nil {
return msg, 0, err
}
n, err = msg.Decode(b)
return msg, n, err
}
// Ack() takes the ack message supplied and updates the status of messages waiting.
func (this *Ackqueue) Ack(msg message.Message) error {
this.mu.Lock()
defer this.mu.Unlock()
switch msg.Type() {
case message.PUBACK, message.PUBREC, message.PUBREL, message.PUBCOMP, message.SUBACK, message.UNSUBACK:
// Check to see if the message w/ the same packet ID is in the queue
i, ok := this.emap[msg.PacketId()]
if ok {
// If message w/ the packet ID exists, update the message state and copy
// the ack message
this.ring[i].State = msg.Type()
ml := msg.Len()
this.ring[i].Ackbuf = make([]byte, ml)
_, err := msg.Encode(this.ring[i].Ackbuf)
if err != nil {
return err
}
//glog.Debugf("Acked: %v", msg)
//} else {
//glog.Debugf("Cannot ack %s message with packet ID %d", msg.Type(), msg.PacketId())
}
case message.PINGRESP:
if this.ping.Mtype == message.PINGREQ {
this.ping.State = message.PINGRESP
}
default:
return errAckMessage
}
return nil
}
func (this *service) processIncoming(msg message.Message) error {
var err error = nil
switch msg := msg.(type) {
case *message.PublishMessage:
// For PUBLISH message, we should figure out what QoS it is and process accordingly
// If QoS == 0, we should just take the next step, no ack required
// If QoS == 1, we should send back PUBACK, then take the next step
// If QoS == 2, we need to put it in the ack queue, send back PUBREC
err = this.processPublish(msg)
case *message.PubackMessage:
// For PUBACK message, it means QoS 1, we should send to ack queue
this.sess.Pub1ack.Ack(msg)
this.processAcked(this.sess.Pub1ack)
case *message.PubrecMessage:
// For PUBREC message, it means QoS 2, we should send to ack queue, and send back PUBREL
if err = this.sess.Pub2out.Ack(msg); err != nil {
break
}
resp := message.NewPubrelMessage()
resp.SetPacketId(msg.PacketId())
_, err = this.writeMessage(resp)
case *message.PubrelMessage:
// For PUBREL message, it means QoS 2, we should send to ack queue, and send back PUBCOMP
if err = this.sess.Pub2in.Ack(msg); err != nil {
break
}
this.processAcked(this.sess.Pub2in)
resp := message.NewPubcompMessage()
resp.SetPacketId(msg.PacketId())
_, err = this.writeMessage(resp)
case *message.PubcompMessage:
// For PUBCOMP message, it means QoS 2, we should send to ack queue
if err = this.sess.Pub2out.Ack(msg); err != nil {
break
}
this.processAcked(this.sess.Pub2out)
case *message.SubscribeMessage:
// For SUBSCRIBE message, we should add subscriber, then send back SUBACK
return this.processSubscribe(msg)
case *message.SubackMessage:
// For SUBACK message, we should send to ack queue
this.sess.Suback.Ack(msg)
this.processAcked(this.sess.Suback)
case *message.UnsubscribeMessage:
// For UNSUBSCRIBE message, we should remove subscriber, then send back UNSUBACK
return this.processUnsubscribe(msg)
case *message.UnsubackMessage:
// For UNSUBACK message, we should send to ack queue
this.sess.Unsuback.Ack(msg)
this.processAcked(this.sess.Unsuback)
case *message.PingreqMessage:
// For PINGREQ message, we should send back PINGRESP
resp := message.NewPingrespMessage()
_, err = this.writeMessage(resp)
case *message.PingrespMessage:
this.sess.Pingack.Ack(msg)
this.processAcked(this.sess.Pingack)
case *message.DisconnectMessage:
// For DISCONNECT message, we should quit
this.sess.Cmsg.SetWillFlag(false)
return errDisconnect
default:
return fmt.Errorf("(%s) invalid message type %s.", this.cid(), msg.Name())
}
if err != nil {
glog.Debugf("(%s) Error processing acked message: %v", this.cid(), err)
}
return err
}
// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
var (
l int = msg.Len()
m, n int
err error
buf []byte
wrap bool
)
if this.out == nil {
return 0, ErrBufferNotReady
}
// This is to serialize writes to the underlying buffer. Multiple goroutines could
// potentially get here because of calling Publish() or Subscribe() or other
// functions that will send messages. For example, if a message is received in
// another connetion, and the message needs to be published to this client, then
// the Publish() function is called, and at the same time, another client could
// do exactly the same thing.
//
// Not an ideal fix though. If possible we should remove mutex and be lockfree.
// Mainly because when there's a large number of goroutines that want to publish
// to this client, then they will all block. However, this will do for now.
//
// FIXME: Try to find a better way than a mutex...if possible.
this.wmu.Lock()
defer this.wmu.Unlock()
buf, wrap, err = this.out.WriteWait(l)
if err != nil {
return 0, err
}
if wrap {
if len(this.outtmp) < l {
this.outtmp = make([]byte, l)
}
n, err = msg.Encode(this.outtmp[0:])
if err != nil {
return 0, err
}
m, err = this.out.Write(this.outtmp[0:n])
if err != nil {
return m, err
}
} else {
n, err = msg.Encode(buf[0:])
if err != nil {
return 0, err
}
m, err = this.out.WriteCommit(n)
if err != nil {
return 0, err
}
}
this.outStat.increment(int64(m))
return m, nil
}
