rclone/vendor/github.com/yunify/qingstor-sdk-go/utils/conn.go

package utils

import (
	"context"
	"net"
	"sync"
	"time"
)

var connPool sync.Pool

type netConn net.Conn
type netDialer net.Dialer

// Dialer is wrapped dialer provided by qingstor go sdk.
//
// We provide this dialer wrapper ReadTimeout & WriteTimeout attributes into connection object.
// This timeout is for individual buffer I/O operation like other language (python, perl... etc),
// so don't bother with SetDeadline or stupid nethttp.Client timeout.
type Dialer struct {
	*net.Dialer
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
}

// NewDialer will create a new dialer.
func NewDialer(connTimeout, readTimeout, writeTimeout time.Duration) *Dialer {
	d := &net.Dialer{
		DualStack: false,
		Timeout:   connTimeout,
	}
	return &Dialer{d, readTimeout, writeTimeout}
}

// Dial connects to the address on the named network.
func (d *Dialer) Dial(network, addr string) (net.Conn, error) {
	c, err := d.Dialer.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	conn := NewConn(c)
	conn.readTimeout = d.ReadTimeout
	conn.writeTimeout = d.WriteTimeout
	return conn, nil
}

// DialContext connects to the address on the named network using
// the provided context.
func (d *Dialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	c, err := d.Dialer.DialContext(ctx, network, addr)
	if err != nil {
		return nil, err
	}
	conn := NewConn(c)
	conn.readTimeout = d.ReadTimeout
	conn.writeTimeout = d.WriteTimeout
	return conn, nil
}

// Conn is a generic stream-oriented network connection.
type Conn struct {
	netConn
	readTimeout  time.Duration
	writeTimeout time.Duration
	timeoutFunc  func() bool
}

// NewConn will create a new conn.
func NewConn(c netConn) *Conn {
	conn, ok := c.(*Conn)
	if ok {
		return conn
	}
	conn, ok = connPool.Get().(*Conn)
	if !ok {
		conn = new(Conn)
	}
	conn.netConn = c
	return conn
}

// SetReadTimeout will set the conn's read timeout.
func (c *Conn) SetReadTimeout(d time.Duration) {
	c.readTimeout = d
}

// SetWriteTimeout will set the conn's write timeout.
func (c *Conn) SetWriteTimeout(d time.Duration) {
	c.writeTimeout = d
}

// Read will read from the conn.
func (c Conn) Read(buf []byte) (n int, err error) {
	if c.readTimeout > 0 {
		c.SetDeadline(time.Now().Add(c.readTimeout))
	}
	n, err = c.netConn.Read(buf)
	if c.readTimeout > 0 {
		c.SetDeadline(time.Time{}) // clear timeout
	}
	return
}

// Write will write into the conn.
func (c Conn) Write(buf []byte) (n int, err error) {
	if c.writeTimeout > 0 {
		c.SetDeadline(time.Now().Add(c.writeTimeout))
	}
	n, err = c.netConn.Write(buf)
	if c.writeTimeout > 0 {
		c.SetDeadline(time.Time{})
	}
	return
}

// Close will close the conn.
func (c Conn) Close() (err error) {
	if c.netConn == nil {
		return nil
	}
	err = c.netConn.Close()
	connPool.Put(c)
	c.netConn = nil
	c.readTimeout = 0
	c.writeTimeout = 0
	return
}

// IsTimeoutError will check whether the err is a timeout error.
func IsTimeoutError(err error) bool {
	e, ok := err.(net.Error)
	if ok {
		return e.Timeout()
	}
	return false
}