implement ContextDial, make connect respect context cancellation

This commit is contained in:
Jeromy
2016-08-16 21:44:43 -07:00
committed by Jakub Sztandera
parent 71193aa1a9
commit 5cf4ab9b06
4 changed files with 152 additions and 59 deletions

View File

@@ -3,6 +3,7 @@
package reuseport
import (
"context"
"net"
"os"
"strconv"
@@ -57,7 +58,7 @@ func socket(family, socktype, protocol int) (fd int, err error) {
return fd, nil
}
func dial(dialer net.Dialer, netw, addr string) (c net.Conn, err error) {
func dial(ctx context.Context, dialer net.Dialer, netw, addr string) (c net.Conn, err error) {
var (
fd int
lfamily int
@@ -135,7 +136,7 @@ func dial(dialer net.Dialer, netw, addr string) (c net.Conn, err error) {
return nil, err
}
if err = connect(fd, remoteSockaddr, deadline); err != nil {
if err = connect(ctx, fd, remoteSockaddr, deadline); err != nil {
syscall.Close(fd)
continue // try again.
}
@@ -301,7 +302,10 @@ func listenUDP(netw, addr string) (c net.Conn, err error) {
}
// this is close to the connect() function inside stdlib/net
func connect(fd int, ra syscall.Sockaddr, deadline time.Time) error {
func connect(ctx context.Context, fd int, ra syscall.Sockaddr, deadline time.Time) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
switch err := syscall.Connect(fd, ra); err {
case syscall.EINPROGRESS, syscall.EALREADY, syscall.EINTR:
case nil, syscall.EISCONN:
@@ -319,6 +323,17 @@ func connect(fd int, ra syscall.Sockaddr, deadline time.Time) error {
}
defer poller.Close()
success := make(chan struct{})
go func() {
<-ctx.Done()
select {
case <-success:
return
default:
poller.Close()
}
}()
for {
if err = poller.WaitWrite(deadline); err != nil {
return err
@@ -344,6 +359,7 @@ func connect(fd int, ra syscall.Sockaddr, deadline time.Time) error {
if !deadline.IsZero() && deadline.Before(time.Now()) {
return errTimeout
}
close(success)
return nil
default:
return err