Merge pull request #68 from fuweid/add-user-on-close-wait

client: add UserOnCloseWait function

client.go

@@ -49,6 +49,7 @@ type Client struct {
 
 	closeOnce       sync.Once
 	userCloseFunc   func()
+	userCloseWaitCh chan struct{}
 
 	errOnce     sync.Once
 	err         error
@@ -82,6 +83,7 @@ func NewClient(conn net.Conn, opts ...ClientOpts) *Client {
 		closed:          cancel,
 		ctx:             ctx,
 		userCloseFunc:   func() {},
+		userCloseWaitCh: make(chan struct{}),
 		interceptor:     defaultClientInterceptor,
 	}
 
@@ -175,6 +177,17 @@ func (c *Client) Close() error {
 	return nil
 }
 
+// UserOnCloseWait is used to blocks untils the user's on-close callback
+// finishes.
+func (c *Client) UserOnCloseWait(ctx context.Context) error {
+	select {
+	case <-c.userCloseWaitCh:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 type message struct {
 	messageHeader
 	p   []byte
@@ -251,6 +264,7 @@ func (c *Client) run() {
 	defer func() {
 		c.conn.Close()
 		c.userCloseFunc()
+		close(c.userCloseWaitCh)
 	}()
 
 	for {

client_test.go

@@ -0,0 +1,70 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package ttrpc
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestUserOnCloseWait(t *testing.T) {
+	var (
+		ctx, cancel    = context.WithDeadline(context.Background(), time.Now().Add(1*time.Minute))
+		server         = mustServer(t)(NewServer())
+		testImpl       = &testingServer{}
+		addr, listener = newTestListener(t)
+	)
+
+	defer cancel()
+	defer listener.Close()
+
+	registerTestingService(server, testImpl)
+
+	go server.Serve(ctx, listener)
+	defer server.Shutdown(ctx)
+
+	var (
+		dataCh          = make(chan string)
+		client, cleanup = newTestClient(t, addr,
+			WithOnClose(func() {
+				dataCh <- time.Now().String()
+			}),
+		)
+
+		tp      testPayload
+		tclient = newTestingClient(client)
+	)
+
+	if _, err := tclient.Test(ctx, &tp); err != nil {
+		t.Fatal(err)
+	}
+
+	cleanup()
+
+	fctx, fcancel := context.WithDeadline(ctx, time.Now().Add(1*time.Second))
+	defer fcancel()
+	if err := client.UserOnCloseWait(fctx); err == nil || err != context.DeadlineExceeded {
+		t.Fatalf("expected error %v, but got %v", context.DeadlineExceeded, err)
+	}
+
+	_ = <-dataCh
+
+	if err := client.UserOnCloseWait(ctx); err != nil {
+		t.Fatalf("expected error nil , but got %v", err)
+	}
+}

server_test.go

@@ -607,12 +607,12 @@ func roundTrip(ctx context.Context, t *testing.T, client *testingClient, value s
 	}
 }
 
-func newTestClient(t testing.TB, addr string) (*Client, func()) {
+func newTestClient(t testing.TB, addr string, opts ...ClientOpts) (*Client, func()) {
 	conn, err := net.Dial("unix", addr)
 	if err != nil {
 		t.Fatal(err)
 	}
-	client := NewClient(conn)
+	client := NewClient(conn, opts...)
 	return client, func() {
 		conn.Close()
 		client.Close()