slight cleanup of dial's timeout

This commit is contained in:
Arceliar 2019-10-21 20:52:16 -05:00
parent 681c8ca6f9
commit ea085663ea

View File

@ -72,19 +72,16 @@ func (d *Dialer) DialByNodeIDandMask(ctx context.Context, nodeID, nodeMask *cryp
return nil, err return nil, err
} }
conn.session.setConn(nil, conn) conn.session.setConn(nil, conn)
var c context.Context
var cancel context.CancelFunc var cancel context.CancelFunc
const timeout = 6 * time.Second if ctx == nil {
if ctx != nil { ctx = context.Background()
c, cancel = context.WithTimeout(ctx, timeout)
} else {
c, cancel = context.WithTimeout(context.Background(), timeout)
} }
ctx, cancel = context.WithTimeout(ctx, 6*time.Second)
defer cancel() defer cancel()
select { select {
case <-conn.session.init: case <-conn.session.init:
return conn, nil return conn, nil
case <-c.Done(): case <-ctx.Done():
conn.Close() conn.Close()
return nil, errors.New("session handshake timeout") return nil, errors.New("session handshake timeout")
} }