shutdown notify channel

Author: DifferentialOrange

connection.go

@@ -144,7 +144,9 @@ func (d defaultLogger) Report(event ConnLogKind, conn *Connection, v ...interfac
 // connection supports server graceful shutdown. In this case, server will
 // wait until all client requests will be finished and client disconnects
 // before going down (server also may go down by timeout). Client reconnect will
-// happen if connection options enable reconnect.
+// happen if connection options enable reconnect. Beware that graceful shutdown
+// event initialization is asynchronous. You may use ShutdownNotify channel
+// in options to track events.
 //
 // More on graceful shutdown: https://www.tarantool.io/en/doc/latest/dev_guide/internals/iproto/graceful_shutdown/
 type Connection struct {
@@ -308,6 +310,8 @@ type Opts struct {
 	// list of protocol features that should be supported by
 	// Tarantool server. By default there are no restrictions.
 	RequiredProtocolInfo ProtocolInfo
+	// ShutdownNotify is a channel which receives Shutdown event values.
+	ShutdownNotify chan<- WatchEvent
 }
 
 // SslOpts is a way to configure ssl transport.
@@ -883,6 +887,15 @@ func (conn *Connection) notify(kind ConnEventKind) {
 	}
 }
 
+func (conn *Connection) shutdownNotify(event WatchEvent) {
+	if conn.opts.ShutdownNotify != nil {
+		select {
+		case conn.opts.ShutdownNotify <- event:
+		default:
+		}
+	}
+}
+
 func (conn *Connection) writer(w *bufio.Writer, c net.Conn) {
 	var shardn uint32
 	var packet smallWBuf
@@ -1740,6 +1753,7 @@ func (conn *Connection) ClientProtocolInfo() ProtocolInfo {
 }
 
 func shutdownEventCallback(event WatchEvent) {
+	event.Conn.shutdownNotify(event)
 	// Receives "true" on server shutdown.
 	// See https://www.tarantool.io/en/doc/latest/dev_guide/internals/iproto/graceful_shutdown/
 	// step 2.

shutdown_test.go

@@ -127,9 +127,15 @@ func TestGracefulShutdown(t *testing.T) {
 	require.Nil(t, err)
 	defer test_helpers.StopTarantoolWithCleanup(inst)
 
+	shtdnNotify := make(chan WatchEvent, 10)
+	shtdnClntOpts.ShutdownNotify = shtdnNotify
+
 	conn = test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
 	defer conn.Close()
 
+	// Wait till Tarantool register out watcher.
+	<-shtdnNotify
+
 	testGracefulShutdown(t, conn, &inst)
 }
 
@@ -143,17 +149,29 @@ func TestGracefulShutdownWithReconnect(t *testing.T) {
 	require.Nil(t, err)
 	defer test_helpers.StopTarantoolWithCleanup(inst)
 
+	shtdnNotify := make(chan WatchEvent, 10)
+	shtdnClntOpts.ShutdownNotify = shtdnNotify
+
 	conn := test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
 	defer conn.Close()
 
+	// Wait till Tarantool register out watcher.
+	<-shtdnNotify
+
 	testGracefulShutdown(t, conn, &inst)
 
+	// Skip shutdown event.
+	<-shtdnNotify
+
 	err = test_helpers.RestartTarantool(&inst)
 	require.Nilf(t, err, "Failed to restart tarantool")
 
 	connected := test_helpers.WaitUntilReconnected(conn, shtdnClntOpts.MaxReconnects, shtdnClntOpts.Reconnect)
 	require.Truef(t, connected, "Reconnect success")
 
+	// Wait till Tarantool register out watcher.
+	<-shtdnNotify
+
 	testGracefulShutdown(t, conn, &inst)
 }
 
@@ -241,9 +259,15 @@ func TestGracefulShutdownRespectsClose(t *testing.T) {
 	require.Nil(t, err)
 	defer test_helpers.StopTarantoolWithCleanup(inst)
 
+	shtdnNotify := make(chan WatchEvent, 10)
+	shtdnClntOpts.ShutdownNotify = shtdnNotify
+
 	conn = test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
 	defer conn.Close()
 
+	// Wait till Tarantool register out watcher.
+	<-shtdnNotify
+
 	// Set a big timeout so it would be easy to differ
 	// if server went down on timeout or after all connections were terminated.
 	serverShutdownTimeout := 60 // in seconds
@@ -311,9 +335,15 @@ func TestGracefulShutdownNotRacesWithRequestReconnect(t *testing.T) {
 	require.Nil(t, err)
 	defer test_helpers.StopTarantoolWithCleanup(inst)
 
+	shtdnNotify := make(chan WatchEvent, 10)
+	shtdnClntOpts.ShutdownNotify = shtdnNotify
+
 	conn = test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
 	defer conn.Close()
 
+	// Wait till Tarantool register out watcher.
+	<-shtdnNotify
+
 	// Set a small timeout so server will shutdown before requesst finishes.
 	serverShutdownTimeout := 1 // in seconds
 	_, err = conn.Call("box.ctl.set_on_shutdown_timeout", []interface{}{serverShutdownTimeout})
@@ -394,6 +424,9 @@ func TestGracefulShutdownCloseConcurrent(t *testing.T) {
 		go func(i int) {
 			defer caseWg.Done()
 
+			// Do not wait till Tarantool register out watcher,
+			// test everything is ok even on async.
+
 			conn := test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
 			defer conn.Close()
 
@@ -463,12 +496,18 @@ func TestGracefulShutdownConcurrent(t *testing.T) {
 	// Create many connections.
 	var ret error
 	for i := 0; i < testConcurrency; i++ {
-		go func(i int) {
+		go func(i int, opts Opts) {
 			defer caseWg.Done()
 
-			conn := test_helpers.ConnectWithValidation(t, shtdnServer, shtdnClntOpts)
+			shtdnNotify := make(chan WatchEvent, 10)
+			opts.ShutdownNotify = shtdnNotify
+
+			conn := test_helpers.ConnectWithValidation(t, shtdnServer, opts)
 			defer conn.Close()
 
+			// Wait till Tarantool register out watcher.
+			<-shtdnNotify
+
 			evalSleep := 1 // in seconds
 			req := NewEvalRequest(evalBody).Args([]interface{}{evalSleep, evalMsg})
 			fut := conn.Do(req)
@@ -482,6 +521,7 @@ func TestGracefulShutdownConcurrent(t *testing.T) {
 				ret = gerr
 			}
 		}(i)
+		}(i, shtdnClntOpts)
 	}
 
 	var sret error