rewise client timeout

client.go

@@ -22,8 +22,7 @@ type Client struct {
 	Timeout time.Duration
 
 	// Dialer allows for customizing the way connections are established.
-	// If set, it overrides the Addr field and uses the provided dialer to
-	// create connections for DNS queries.
+	// If set, Addr and Timeout will be ignore.
 	Dialer Dialer
 }
 
@@ -48,18 +47,18 @@ func (c *Client) exchange(ctx context.Context, req, resp *Message) error {
 		return err
 	}
 
-	_, err = conn.Write(req.Raw)
-	if err != nil {
-		return nil
-	}
-
-	if c.Timeout > 0 {
+	if c.Timeout > 0 && c.Dialer == nil {
 		err = conn.SetDeadline(time.Now().Add(c.Timeout))
 		if err != nil {
 			return err
 		}
 	}
 
+	_, err = conn.Write(req.Raw)
+	if err != nil {
+		return nil
+	}
+
 	resp.Raw = resp.Raw[:cap(resp.Raw)]
 	n, err := conn.Read(resp.Raw)
 	if err != nil {

client_dialer.go

@@ -20,6 +20,10 @@ type UDPDialer struct {
 	// Addr specifies the remote UDP address that the dialer will connect to.
 	Addr *net.UDPAddr
 
+	// Timeout specifies the maximum duration for a query to complete.
+	// If a query exceeds this duration, it will result in a timeout error.
+	Timeout time.Duration
+
 	// MaxConns limits the maximum number of UDP connections that can be created
 	// and reused. Once this limit is reached, no new connections will be made.
 	MaxConns uint16
@@ -43,6 +47,9 @@ func (d *UDPDialer) get() (net.Conn, error) {
 
 	c := d.conns[cheaprandn(uint32(d.MaxConns))]
 	c.mu.Lock()
+	if d.Timeout > 0 {
+		c.SetDeadline(time.Now().Add(d.Timeout))
+	}
 
 	return c, nil
 }