github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/transport/client_fasthttp.go (about)

     1  //go:build !nethttp
     2  
     3  // Package transport provides long-lived http/tcp connections for
     4  // intra-cluster communications (see README for details and usage example).
     5  /*
     6   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     7   */
     8  package transport
     9  
    10  import (
    11  	"io"
    12  	"net"
    13  	"net/http"
    14  	"strconv"
    15  	"time"
    16  
    17  	"github.com/NVIDIA/aistore/api/apc"
    18  	"github.com/NVIDIA/aistore/cmn"
    19  	"github.com/NVIDIA/aistore/cmn/cos"
    20  	"github.com/NVIDIA/aistore/cmn/nlog"
    21  	"github.com/valyala/fasthttp"
    22  )
    23  
    24  const ua = "aisnode/streams"
    25  
    26  type Client interface {
    27  	Do(req *fasthttp.Request, resp *fasthttp.Response) error
    28  }
    29  
    30  func whichClient() string { return "fasthttp" }
    31  
    32  // overriding fasthttp default `const DefaultDialTimeout = 3 * time.Second`
    33  func dialTimeout(addr string) (net.Conn, error) {
    34  	return fasthttp.DialTimeout(addr, 10*time.Second)
    35  }
    36  
    37  // intra-cluster networking: fasthttp client
    38  func NewIntraDataClient() Client {
    39  	config := cmn.GCO.Get()
    40  
    41  	// compare with ais/httpcommon.go
    42  	wbuf, rbuf := config.Net.HTTP.WriteBufferSize, config.Net.HTTP.ReadBufferSize
    43  	if wbuf == 0 {
    44  		wbuf = cmn.DefaultWriteBufferSize // fasthttp uses 4KB
    45  	}
    46  	if rbuf == 0 {
    47  		rbuf = cmn.DefaultReadBufferSize // ditto
    48  	}
    49  	cl := &fasthttp.Client{
    50  		Dial:            dialTimeout,
    51  		ReadBufferSize:  rbuf,
    52  		WriteBufferSize: wbuf,
    53  	}
    54  	if config.Net.HTTP.UseHTTPS {
    55  		tlsConfig, err := cmn.NewTLS(config.Net.HTTP.ToTLS())
    56  		if err != nil {
    57  			cos.ExitLog(err)
    58  		}
    59  		cl.TLSConfig = tlsConfig
    60  	}
    61  	return cl
    62  }
    63  
    64  func (s *streamBase) do(body io.Reader) (err error) {
    65  	// init request & response
    66  	req, resp := fasthttp.AcquireRequest(), fasthttp.AcquireResponse()
    67  	req.Header.SetMethod(http.MethodPut)
    68  	req.SetRequestURI(s.dstURL)
    69  	req.SetBodyStream(body, -1)
    70  	if s.streamer.compressed() {
    71  		req.Header.Set(apc.HdrCompress, apc.LZ4Compression)
    72  	}
    73  	req.Header.Set(apc.HdrSessID, strconv.FormatInt(s.sessID, 10))
    74  	req.Header.Set(cos.HdrUserAgent, ua)
    75  	// do
    76  	err = s.client.Do(req, resp)
    77  	if err != nil {
    78  		if verbose {
    79  			nlog.Errorf("%s: Error [%v]", s, err)
    80  		}
    81  		return
    82  	}
    83  	// handle response & cleanup
    84  	resp.BodyWriteTo(io.Discard)
    85  	fasthttp.ReleaseRequest(req)
    86  	fasthttp.ReleaseResponse(resp)
    87  	if s.streamer.compressed() {
    88  		s.streamer.resetCompression()
    89  	}
    90  	return
    91  }