github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/request/options.go (about)

     1  package request
     2  
     3  import (
     4  	"context"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
     6  	"net/http"
     7  	"net/url"
     8  	"strings"
     9  	"time"
    10  )
    11  
    12  // Option 发送请求的额外设置
    13  type Option interface {
    14  	apply(*options)
    15  }
    16  
    17  type options struct {
    18  	timeout         time.Duration
    19  	header          http.Header
    20  	sign            auth.Auth
    21  	signTTL         int64
    22  	ctx             context.Context
    23  	contentLength   int64
    24  	masterMeta      bool
    25  	endpoint        *url.URL
    26  	slaveNodeID     string
    27  	tpsLimiterToken string
    28  	tps             float64
    29  	tpsBurst        int
    30  }
    31  
    32  type optionFunc func(*options)
    33  
    34  func (f optionFunc) apply(o *options) {
    35  	f(o)
    36  }
    37  
    38  func newDefaultOption() *options {
    39  	return &options{
    40  		header:        http.Header{},
    41  		timeout:       time.Duration(30) * time.Second,
    42  		contentLength: -1,
    43  		ctx:           context.Background(),
    44  	}
    45  }
    46  
    47  func (o *options) clone() options {
    48  	newOptions := *o
    49  	newOptions.header = o.header.Clone()
    50  	return newOptions
    51  }
    52  
    53  // WithTimeout 设置请求超时
    54  func WithTimeout(t time.Duration) Option {
    55  	return optionFunc(func(o *options) {
    56  		o.timeout = t
    57  	})
    58  }
    59  
    60  // WithContext 设置请求上下文
    61  func WithContext(c context.Context) Option {
    62  	return optionFunc(func(o *options) {
    63  		o.ctx = c
    64  	})
    65  }
    66  
    67  // WithCredential 对请求进行签名
    68  func WithCredential(instance auth.Auth, ttl int64) Option {
    69  	return optionFunc(func(o *options) {
    70  		o.sign = instance
    71  		o.signTTL = ttl
    72  	})
    73  }
    74  
    75  // WithHeader 设置请求Header
    76  func WithHeader(header http.Header) Option {
    77  	return optionFunc(func(o *options) {
    78  		for k, v := range header {
    79  			o.header[k] = v
    80  		}
    81  	})
    82  }
    83  
    84  // WithoutHeader 设置清除请求Header
    85  func WithoutHeader(header []string) Option {
    86  	return optionFunc(func(o *options) {
    87  		for _, v := range header {
    88  			delete(o.header, v)
    89  		}
    90  
    91  	})
    92  }
    93  
    94  // WithContentLength 设置请求大小
    95  func WithContentLength(s int64) Option {
    96  	return optionFunc(func(o *options) {
    97  		o.contentLength = s
    98  	})
    99  }
   100  
   101  // WithMasterMeta 请求时携带主机信息
   102  func WithMasterMeta() Option {
   103  	return optionFunc(func(o *options) {
   104  		o.masterMeta = true
   105  	})
   106  }
   107  
   108  // WithSlaveMeta 请求时携带从机信息
   109  func WithSlaveMeta(s string) Option {
   110  	return optionFunc(func(o *options) {
   111  		o.slaveNodeID = s
   112  	})
   113  }
   114  
   115  // Endpoint 使用同一的请求Endpoint
   116  func WithEndpoint(endpoint string) Option {
   117  	if !strings.HasSuffix(endpoint, "/") {
   118  		endpoint += "/"
   119  	}
   120  
   121  	endpointURL, _ := url.Parse(endpoint)
   122  	return optionFunc(func(o *options) {
   123  		o.endpoint = endpointURL
   124  	})
   125  }
   126  
   127  // WithTPSLimit 请求时使用全局流量限制
   128  func WithTPSLimit(token string, tps float64, burst int) Option {
   129  	return optionFunc(func(o *options) {
   130  		o.tpsLimiterToken = token
   131  		o.tps = tps
   132  		if burst < 1 {
   133  			burst = 1
   134  		}
   135  		o.tpsBurst = burst
   136  	})
   137  }