Make onclose an option.

Signed-off-by: Lantao Liu <lantaol@google.com>

client.go

@@ -49,7 +49,15 @@ type Client struct {
 	err       error
 }
 
-func NewClient(conn net.Conn) *Client {
+type ClientOpts func(c *Client)
+
+func WithOnClose(onClose func()) ClientOpts {
+	return func(c *Client) {
+		c.closeFunc = onClose
+	}
+}
+
+func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
 	c := &Client{
 		codec:     codec{},
 		conn:      conn,
@@ -60,6 +68,10 @@ func NewClient(conn net.Conn) *Client {
 		closeFunc: func() {},
 	}
 
+	for _, o := range opts {
+		o(c)
+	}
+
 	go c.run()
 	return c
 }
@@ -141,11 +153,6 @@ 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