Merge pull request #121 from sime1/context

add DialPipeContext

pipe.go

@@ -3,6 +3,7 @@
 package winio
 
 import (
+	"context"
 	"errors"
 	"io"
 	"net"
@@ -137,9 +138,30 @@ func (s pipeAddress) String() string {
 	return string(s)
 }
 
+// tryDialPipe attempts to dial the pipe at `path` until `ctx` cancellation or timeout.
+func tryDialPipe(ctx context.Context, path *string) (syscall.Handle, error) {
+	for {
+		select {
+		case <-ctx.Done():
+			return syscall.Handle(0), ctx.Err()
+		default:
+			h, err := createFile(*path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, 0, nil, syscall.OPEN_EXISTING, syscall.FILE_FLAG_OVERLAPPED|cSECURITY_SQOS_PRESENT|cSECURITY_ANONYMOUS, 0)
+			if err == nil {
+				return h, nil
+			}
+			if err != cERROR_PIPE_BUSY {
+				return h, &os.PathError{Err: err, Op: "open", Path: *path}
+			}
+			// Wait 10 msec and try again. This is a rather simplistic
+			// view, as we always try each 10 milliseconds.
+			time.Sleep(time.Millisecond * 10)
+		}
+	}
+}
+
 // DialPipe connects to a named pipe by path, timing out if the connection
 // takes longer than the specified duration. If timeout is nil, then we use
-// a default timeout of 5 seconds.  (We do not use WaitNamedPipe.)
+// a default timeout of 2 seconds.  (We do not use WaitNamedPipe.)
 func DialPipe(path string, timeout *time.Duration) (net.Conn, error) {
 	var absTimeout time.Time
 	if timeout != nil {
@@ -147,23 +169,22 @@ func DialPipe(path string, timeout *time.Duration) (net.Conn, error) {
 	} else {
 		absTimeout = time.Now().Add(time.Second * 2)
 	}
+	ctx, _ := context.WithDeadline(context.Background(), absTimeout)
+	conn, err := DialPipeContext(ctx, path)
+	if err == context.DeadlineExceeded {
+		return nil, ErrTimeout
+	}
+	return conn, err
+}
+
+// DialPipeContext attempts to connect to a named pipe by `path` until `ctx`
+// cancellation or timeout.
+func DialPipeContext(ctx context.Context, path string) (net.Conn, error) {
 	var err error
 	var h syscall.Handle
-	for {
-		h, err = createFile(path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, 0, nil, syscall.OPEN_EXISTING, syscall.FILE_FLAG_OVERLAPPED|cSECURITY_SQOS_PRESENT|cSECURITY_ANONYMOUS, 0)
-		if err != cERROR_PIPE_BUSY {
-			break
-		}
-		if time.Now().After(absTimeout) {
-			return nil, ErrTimeout
-		}
-
-		// Wait 10 msec and try again. This is a rather simplistic
-		// view, as we always try each 10 milliseconds.
-		time.Sleep(time.Millisecond * 10)
-	}
+	h, err = tryDialPipe(ctx, &path)
 	if err != nil {
-		return nil, &os.PathError{Op: "open", Path: path, Err: err}
+		return nil, err
 	}
 
 	var flags uint32

pipe_test.go

@@ -3,6 +3,7 @@ package winio
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"io"
 	"net"
 	"os"
@@ -37,6 +38,40 @@ func TestDialListenerTimesOut(t *testing.T) {
 	}
 }
 
+func TestDialContextListenerTimesOut(t *testing.T) {
+	l, err := ListenPipe(testPipeName, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+	var d = time.Duration(10 * time.Millisecond)
+	ctx, _ := context.WithTimeout(context.Background(), d)
+	_, err = DialPipeContext(ctx, testPipeName)
+	if err != context.DeadlineExceeded {
+		t.Fatalf("expected context.DeadlineExceeded, got %v", err)
+	}
+}
+
+func TestDialListenerGetsCancelled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	l, err := ListenPipe(testPipeName, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ch := make(chan error)
+	defer l.Close()
+	go func(ctx context.Context, ch chan error) {
+		_, err := DialPipeContext(ctx, testPipeName)
+		ch <- err
+	}(ctx, ch)
+	time.Sleep(time.Millisecond * 30)
+	cancel()
+	err = <-ch
+	if err != context.Canceled {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+}
+
 func TestDialAccessDeniedWithRestrictedSD(t *testing.T) {
 	c := PipeConfig{
 		SecurityDescriptor: "D:P(A;;0x1200FF;;;WD)",