Add onclose func

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>

client.go

@@ -27,6 +27,7 @@ type Client struct {
 
 	closed    chan struct{}
 	closeOnce sync.Once
+	closeFunc func()
 	done      chan struct{}
 	err       error
 }
@@ -39,6 +40,7 @@ func NewClient(conn net.Conn) *Client {
 		calls:     make(chan *callRequest),
 		closed:    make(chan struct{}),
 		done:      make(chan struct{}),
+		closeFunc: func() {},
 	}
 
 	go c.run()
@@ -113,6 +115,11 @@ func (c *Client) Close() error {
 	return nil
 }
 
+// OnClose allows a close func to be called when the server is closed
+func (c *Client) OnClose(closer func()) {
+	c.closeFunc = closer
+}
+
 type message struct {
 	messageHeader
 	p   []byte
@@ -158,6 +165,7 @@ func (c *Client) run() {
 
 	defer c.conn.Close()
 	defer close(c.done)
+	defer c.closeFunc()
 
 	for {
 		select {