conditionally announce message on Publish (enabling server to retry) …

…or hold until Pubrel (having certainty the broker won't redeliver).

Add tests for both behaviours.

client/client.go

@@ -113,6 +113,7 @@ type Client struct {
 	config        *Config
 	conn          transport.Conn
 	clean         bool
+	earlyCallback bool
 	keepAlive     time.Duration
 	tracker       *Tracker
 	futureStore   *future.Store
@@ -206,6 +207,8 @@ func (c *Client) Connect(config *Config) (ConnectFuture, error) {
 		}
 	}
 
+	c.earlyCallback = config.AlwaysAnnounceOnPublish
+
 	// allocate packet
 	connect := packet.NewConnect()
 	connect.ClientID = config.ClientID
@@ -611,15 +614,14 @@ func (c *Client) processUnsuback(unsuback *packet.Unsuback) error {
 // handle an incoming Publish packet
 func (c *Client) processPublish(publish *packet.Publish) error {
 	// call callback for unacknowledged and directly acknowledged messages
-	if publish.Message.QOS <= 1 {
+	if publish.Message.QOS <= 1 || c.earlyCallback {
 		if c.Callback != nil {
 			err := c.Callback(&publish.Message, nil)
 			if err != nil {
 				return c.die(err, true)
 			}
 		}
 	}
-
 	// handle qos 1 flow
 	if publish.Message.QOS == 1 {
 		// prepare puback packet
@@ -714,7 +716,7 @@ func (c *Client) processPubrel(id packet.ID) error {
 	}
 
 	// call callback
-	if c.Callback != nil {
+	if c.Callback != nil && !c.earlyCallback {
 		err = c.Callback(&publish.Message, nil)
 		if err != nil {
 			return c.die(err, true)

client/config.go

@@ -41,6 +41,14 @@ type Config struct {
 	// MaxWriteDelay defines the maximum allowed delay when flushing the
 	// underlying buffered writer.
 	MaxWriteDelay time.Duration
+
+	// AlwaysAnnounceOnPublish defines when the message callback handler is called.
+	// QOS 0,1: callback always occurs on reception of Publish
+	// QOS 2 and AlwaysAnnounceOnPublish == false: callback occurs on reception of PubRel and returning and error in
+	//   the callback will close the connection before PubComp is sent.
+	// QOS 2 and AlwaysAnnounceOnPublish == true: callback occurs on reception of Publish and returning an error in
+	//   the callback will close the connection, preventing PubRec being sent, thus ensuring redelivery of publish.
+	AlwaysAnnounceOnPublish bool
 }
 
 // NewConfig creates a new Config using the specified URL.