[ADDED] Allow client to choose whether to execute callbacks on Close()

nats.go

@@ -347,6 +347,11 @@ type Options struct {
 	// UseOldRequestStyle forces the old method of Requests that utilize
 	// a new Inbox and a new Subscription for each request.
 	UseOldRequestStyle bool
+
+	// NoCallbacksAfterClientClose allows preventing the invocation of
+	// callbacks after Close() is called. Client won't receive notifications
+	// when Close is invoked by user code. Default is to invoke the callbacks.
+	NoCallbacksAfterClientClose bool
 }
 
 const (
@@ -883,6 +888,16 @@ func UseOldRequestStyle() Option {
 	}
 }
 
+// NoCallbacksAfterClientClose is an Option to disable callbacks when user code
+// calls Close(). If close is initiated by any other condition, callbacks
+// if any will be invoked.
+func NoCallbacksAfterClientClose() Option {
+	return func(o *Options) error {
+		o.NoCallbacksAfterClientClose = true
+		return nil
+	}
+}
+
 // Handler processing
 
 // SetDisconnectHandler will set the disconnect event handler.
@@ -1929,7 +1944,7 @@ func (nc *Conn) doReconnect(err error) {
 		nc.err = ErrNoServers
 	}
 	nc.mu.Unlock()
-	nc.Close()
+	nc.close(CLOSED, true, nil)
 }
 
 // processOpErr handles errors from reading or parsing the protocol.
@@ -1964,7 +1979,7 @@ func (nc *Conn) processOpErr(err error) {
 	nc.status = DISCONNECTED
 	nc.err = err
 	nc.mu.Unlock()
-	nc.Close()
+	nc.close(CLOSED, true, nil)
 }
 
 // dispatch is responsible for calling any async callbacks
@@ -2496,7 +2511,7 @@ func (nc *Conn) processErr(ie string) {
 		nc.mu.Unlock()
 	}
 	if close {
-		nc.Close()
+		nc.close(CLOSED, true, nil)
 	}
 }
 
@@ -3691,7 +3706,7 @@ func (nc *Conn) close(status Status, doCBs bool, err error) {
 // all blocking calls, such as Flush() and NextMsg()
 func (nc *Conn) Close() {
 	if nc != nil {
-		nc.close(CLOSED, true, nil)
+		nc.close(CLOSED, !nc.Opts.NoCallbacksAfterClientClose, nil)
 	}
 }
 
@@ -3770,12 +3785,12 @@ func (nc *Conn) drainConnection() {
 	err := nc.Flush()
 	if err != nil {
 		pushErr(err)
-		nc.Close()
+		nc.close(CLOSED, true, nil)
 		return
 	}
 
 	// Move to closed state.
-	nc.Close()
+	nc.close(CLOSED, true, nil)
 }
 
 // Drain will put a connection into a drain state. All subscriptions will

test/reconnect_test.go

@@ -722,3 +722,66 @@ func TestReconnectTLSHostNoIP(t *testing.T) {
 		t.Fatalf("ReconnectedCB should have been triggered: %v", nc.LastError())
 	}
 }
+
+func TestConnCloseNoCallback(t *testing.T) {
+	ts := startReconnectServer(t)
+	defer ts.Shutdown()
+
+	cch := make(chan bool)
+	opts := reconnectOpts
+	opts.ClosedCB = func(_ *nats.Conn) {
+		cch <- true
+	}
+	opts.NoCallbacksAfterClientClose = true
+
+	nc, err := opts.Connect()
+	if err != nil {
+		t.Fatalf("Should have connected ok: %v", err)
+	}
+	defer nc.Close()
+
+	nc.Flush()
+	// Close the connection, we don't expect to get a notification
+	nc.Close()
+	// Shutdown the server
+	ts.Shutdown()
+
+	// Even on Windows (where a createConn takes more than a second)
+	// we should be able to break the reconnect loop with the following
+	// timeout.
+	if err := WaitTime(cch, 3*time.Second); err != nil {
+		// yay no callback
+	} else {
+		t.Fatal("Got a closed callback, but shouldn't have")
+	}
+}
+
+func TestConnCloseNoCallbackFromOptionsFunc(t *testing.T) {
+	ts := startReconnectServer(t)
+	defer ts.Shutdown()
+
+	cch := make(chan bool)
+	nc, err := nats.Connect(reconnectOpts.Url, nats.NoCallbacksAfterClientClose(),
+		nats.ClosedHandler(func(_ *nats.Conn) {
+			cch <- true
+		}))
+	if err != nil {
+		t.Fatalf("Should have connected ok: %v", err)
+	}
+	defer nc.Close()
+
+	nc.Flush()
+	// Close the connection, we don't expect to get a notification
+	nc.Close()
+	// Shutdown the server
+	ts.Shutdown()
+
+	// Even on Windows (where a createConn takes more than a second)
+	// we should be able to break the reconnect loop with the following
+	// timeout.
+	if err := WaitTime(cch, 3*time.Second); err != nil {
+		// yay no callback
+	} else {
+		t.Fatal("Got a closed callback, but shouldn't have")
+	}
+}