api: add Connection.CloseGraceful()

CloseGraceful closes Connection gracefully. Unlike Connection.Close()
it waits for all requests to complete.

Part of #257

Author: oleg-jukovec

CHANGELOG.md

@@ -10,6 +10,9 @@ Versioning](http://semver.org/spec/v2.0.0.html) except to the first release.
 
 ### Added
 
+- Connection.CloseGraceful() unlike Connection.Close() waits for all
+  requests to complete (#257)
+
 ### Changed
 
 ### Fixed

connection.go

@@ -462,6 +462,13 @@ func (conn *Connection) Close() error {
 	return conn.closeConnection(err, true)
 }
 
+// CloseGraceful closes Connection gracefully. It waits for all requests to
+// complete.
+// After this method called, there is no way to reopen this Connection.
+func (conn *Connection) CloseGraceful() error {
+	return conn.shutdown(true)
+}
+
 // Addr returns a configured address of Tarantool socket.
 func (conn *Connection) Addr() string {
 	return conn.addr
@@ -1532,17 +1539,27 @@ func shutdownEventCallback(event WatchEvent) {
 	// step 2.
 	val, ok := event.Value.(bool)
 	if ok && val {
-		go event.Conn.shutdown()
+		go event.Conn.shutdown(false)
 	}
 }
 
-func (conn *Connection) shutdown() {
+func (conn *Connection) shutdown(forever bool) error {
 	// Forbid state changes.
 	conn.mutex.Lock()
 	defer conn.mutex.Unlock()
 
 	if !atomic.CompareAndSwapUint32(&conn.state, connConnected, connShutdown) {
-		return
+		if forever {
+			err := ClientError{ErrConnectionClosed, "connection closed by client"}
+			return conn.closeConnection(err, true)
+		}
+		return nil
+	}
+
+	if forever {
+		// We don't want to reconnect any more.
+		conn.opts.Reconnect = 0
+		conn.opts.MaxReconnects = 0
 	}
 
 	conn.cond.Broadcast()
@@ -1551,7 +1568,7 @@ func (conn *Connection) shutdown() {
 	c := conn.c
 	for {
 		if (atomic.LoadUint32(&conn.state) != connShutdown) || (c != conn.c) {
-			return
+			return nil
 		}
 		if atomic.LoadInt64(&conn.requestCnt) == 0 {
 			break
@@ -1563,14 +1580,19 @@ func (conn *Connection) shutdown() {
 		conn.cond.Wait()
 	}
 
-	// Start to reconnect based on common rules, same as in net.box.
-	// Reconnect also closes the connection: server waits until all
-	// subscribed connections are terminated.
-	// See https://www.tarantool.io/en/doc/latest/dev_guide/internals/iproto/graceful_shutdown/
-	// step 3.
-	conn.reconnectImpl(
-		ClientError{
+	if forever {
+		err := ClientError{ErrConnectionClosed, "connection closed by client"}
+		return conn.closeConnection(err, true)
+	} else {
+		// Start to reconnect based on common rules, same as in net.box.
+		// Reconnect also closes the connection: server waits until all
+		// subscribed connections are terminated.
+		// See https://www.tarantool.io/en/doc/latest/dev_guide/internals/iproto/graceful_shutdown/
+		// step 3.
+		conn.reconnectImpl(ClientError{
 			ErrConnectionClosed,
 			"connection closed after server shutdown",
 		}, conn.c)
+		return nil
+	}
 }

example_test.go

@@ -1115,3 +1115,39 @@ func ExamplePingRequest_Context() {
 	// Ping Resp <nil>
 	// Ping Error context is done
 }
+
+// ExampleConnection_CloseGraceful_force demonstrates how to force close
+// a connection with graceful close in progress after a while.
+func ExampleConnection_CloseGraceful_force() {
+	conn := example_connect(opts)
+
+	eval := `local fiber = require('fiber')
+	local time = ...
+	fiber.sleep(time)
+`
+	req := tarantool.NewEvalRequest(eval).Args([]interface{}{10})
+	fut := conn.Do(req)
+
+	done := make(chan struct{})
+	go func() {
+		conn.CloseGraceful()
+		fmt.Println("Connection.CloseGraceful() done!")
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		fmt.Println("Force Connection.Close()!")
+		conn.Close()
+	}
+	<-done
+
+	fmt.Println("Result:")
+	fmt.Println(fut.Get())
+	// Output:
+	// Force Connection.Close()!
+	// Connection.CloseGraceful() done!
+	// Result:
+	// <nil> connection closed by client (0x4001)
+}

shutdown_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	. "github.com/tarantool/go-tarantool"
 	"github.com/tarantool/go-tarantool/test_helpers"
@@ -143,6 +144,48 @@ func TestGracefulShutdown(t *testing.T) {
 	testGracefulShutdown(t, conn, &inst)
 }
 
+func TestCloseGraceful(t *testing.T) {
+	opts := Opts{
+		User:    shtdnClntOpts.User,
+		Pass:    shtdnClntOpts.Pass,
+		Timeout: shtdnClntOpts.Timeout,
+	}
+
+	inst, err := test_helpers.StartTarantool(shtdnSrvOpts)
+	require.Nil(t, err)
+	defer test_helpers.StopTarantoolWithCleanup(inst)
+
+	conn := test_helpers.ConnectWithValidation(t, shtdnServer, opts)
+	defer conn.Close()
+
+	// Send request with sleep.
+	evalSleep := 3 // In seconds.
+	require.Lessf(t,
+		time.Duration(evalSleep)*time.Second,
+		shtdnClntOpts.Timeout,
+		"test request won't be failed by timeout")
+
+	req := NewEvalRequest(evalBody).Args([]interface{}{evalSleep, evalMsg})
+	fut := conn.Do(req)
+
+	go func() {
+		// CloseGraceful closes the connection gracefully.
+		conn.CloseGraceful()
+		// Connection is closed.
+		assert.Equal(t, true, conn.ClosedNow())
+	}()
+
+	// Check that a request rejected if graceful shutdown in progress.
+	time.Sleep((time.Duration(evalSleep) * time.Second) / 2)
+	_, err = conn.Do(NewPingRequest()).Get()
+	assert.ErrorContains(t, err, "server shutdown in progress")
+
+	// Check that a previous request was successful.
+	resp, err := fut.Get()
+	assert.Nilf(t, err, "sleep request no error")
+	assert.NotNilf(t, resp, "sleep response exists")
+}
+
 func TestGracefulShutdownWithReconnect(t *testing.T) {
 	test_helpers.SkipIfWatchersUnsupported(t)