From 7ba6d478f0a146b659e3e72f82a9f5e1bf76f843 Mon Sep 17 00:00:00 2001
From: Robbie Trencheny <me@robbiet.us>
Date: Mon, 30 Jan 2017 14:14:07 -0800
Subject: [PATCH] Add a handler for message successfully sent

---
 melody.go  | 10 +++++++++-
 session.go |  4 +++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/melody.go b/melody.go
index 0ec990a..e0c20bc 100644
--- a/melody.go
+++ b/melody.go
@@ -1,9 +1,10 @@
 package melody
 
 import (
-	"github.com/gorilla/websocket"
 	"net/http"
 	"sync"
+
+	"github.com/gorilla/websocket"
 )
 
 type handleMessageFunc func(*Session, []byte)
@@ -17,6 +18,7 @@ type Melody struct {
 	Upgrader             *websocket.Upgrader
 	messageHandler       handleMessageFunc
 	messageHandlerBinary handleMessageFunc
+	messageSentHandler   handleMessageFunc
 	errorHandler         handleErrorFunc
 	connectHandler       handleSessionFunc
 	disconnectHandler    handleSessionFunc
@@ -40,6 +42,7 @@ func New() *Melody {
 		Upgrader:             upgrader,
 		messageHandler:       func(*Session, []byte) {},
 		messageHandlerBinary: func(*Session, []byte) {},
+		messageSentHandler:   func(*Session, []byte) {},
 		errorHandler:         func(*Session, error) {},
 		connectHandler:       func(*Session) {},
 		disconnectHandler:    func(*Session) {},
@@ -73,6 +76,11 @@ func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) {
 	m.messageHandlerBinary = fn
 }
 
+// HandleSentMessage fires fn when a message is successfully sent.
+func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) {
+	m.messageSentHandler = fn
+}
+
 // HandleError fires fn when a session has an error.
 func (m *Melody) HandleError(fn func(*Session, error)) {
 	m.errorHandler = fn
diff --git a/session.go b/session.go
index 1c02769..745fa90 100644
--- a/session.go
+++ b/session.go
@@ -2,10 +2,11 @@ package melody
 
 import (
 	"errors"
-	"github.com/gorilla/websocket"
 	"net/http"
 	"sync"
 	"time"
+
+	"github.com/gorilla/websocket"
 )
 
 // Session is wrapper around websocket connections.
@@ -71,6 +72,7 @@ loop:
 				s.melody.errorHandler(s, err)
 				break loop
 			}
+			s.melody.messageSentHandler(s, msg.msg)
 		case <-ticker.C:
 			s.ping()
 		}