github.com/kaydxh/golang@v0.0.131/pkg/file-transfer/file.transfer.go (about)

     1  /*
     2   *Copyright (c) 2022, kaydxh
     3   *
     4   *Permission is hereby granted, free of charge, to any person obtaining a copy
     5   *of this software and associated documentation files (the "Software"), to deal
     6   *in the Software without restriction, including without limitation the rights
     7   *to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8   *copies of the Software, and to permit persons to whom the Software is
     9   *furnished to do so, subject to the following conditions:
    10   *
    11   *The above copyright notice and this permission notice shall be included in all
    12   *copies or substantial portions of the Software.
    13   *
    14   *THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15   *IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16   *FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17   *AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18   *LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19   *OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    20   *SOFTWARE.
    21   */
    22  package filetransfer
    23  
    24  import (
    25  	"context"
    26  	"math/rand"
    27  	"time"
    28  
    29  	http_ "github.com/kaydxh/golang/go/net/http"
    30  	time_ "github.com/kaydxh/golang/go/time"
    31  	logs_ "github.com/kaydxh/golang/pkg/logs"
    32  	"go.opentelemetry.io/otel"
    33  )
    34  
    35  type FileTransferOptions struct {
    36  	// 0 means no timeout
    37  	downloadTimeout time.Duration
    38  	uploadTimeout   time.Duration
    39  	loadBalanceMode Ft_LoadBalanceMode
    40  	retryTimes      int
    41  	// retry interval, 0 means retry immediately
    42  	retryInterval time.Duration
    43  
    44  	proxies []*Ft_Proxy
    45  }
    46  
    47  func defaultFileTransferOptions() FileTransferOptions {
    48  	return FileTransferOptions{}
    49  }
    50  
    51  type FileTransfer struct {
    52  	opts FileTransferOptions
    53  }
    54  
    55  func NewFileTransfer(opts ...FileTransferOption) *FileTransfer {
    56  	ft := &FileTransfer{}
    57  	ft.ApplyOptions(opts...)
    58  
    59  	return ft
    60  }
    61  
    62  func (f *FileTransfer) getProxy() *Ft_Proxy {
    63  	proxies := f.opts.proxies
    64  	if len(proxies) == 0 {
    65  		return &Ft_Proxy{}
    66  	}
    67  
    68  	switch f.opts.loadBalanceMode {
    69  	case Ft_load_balance_mode_random:
    70  		return proxies[rand.Intn(len(proxies))]
    71  	default:
    72  		return proxies[0]
    73  	}
    74  }
    75  
    76  // short connection
    77  func (f *FileTransfer) Download(ctx context.Context, downloadUrl string) (data []byte, err error) {
    78  	spanName := "Download"
    79  	ctx, span := otel.Tracer("").Start(ctx, spanName)
    80  	defer span.End()
    81  
    82  	logger := logs_.GetLogger(ctx)
    83  	logger = logger.WithField("trace_id", span.SpanContext().TraceID()).WithField("span_id", span.SpanContext().SpanID()).WithField("download_url", downloadUrl)
    84  
    85  	proxy := f.getProxy()
    86  
    87  	opts := []http_.ClientOption{http_.WithDisableKeepAlives(true)}
    88  	if proxy.TargetHost != "" {
    89  		opts = append(opts, http_.WithTargetHost(proxy.TargetHost))
    90  	} else {
    91  
    92  		if proxy.ProxyUrl != "" {
    93  			opts = append(opts, http_.WithProxyURL(proxy.ProxyUrl))
    94  		}
    95  		if proxy.ProxyHost != "" {
    96  			opts = append(opts, http_.WithProxyHost(proxy.ProxyHost))
    97  		}
    98  	}
    99  	opts = append(opts, http_.WithTimeout(f.opts.downloadTimeout))
   100  
   101  	err = time_.RetryWithContext(ctx, func(ctx context.Context) error {
   102  		client, err := http_.NewClient(opts...)
   103  		if err != nil {
   104  			logger.WithError(err).Errorf("new http client err: %v", err)
   105  			return err
   106  		}
   107  
   108  		data, err = client.Get(ctx, downloadUrl)
   109  		if err != nil {
   110  			logger.WithError(err).Errorf("http client get err: %v", err)
   111  			return err
   112  		}
   113  		return nil
   114  
   115  	}, f.opts.retryInterval, f.opts.retryTimes)
   116  
   117  	return data, err
   118  }
   119  
   120  // short connection
   121  func (f *FileTransfer) Upload(ctx context.Context, uploadUrl string, body []byte) (data []byte, err error) {
   122  	spanName := "Upload"
   123  	ctx, span := otel.Tracer("").Start(ctx, spanName)
   124  	defer span.End()
   125  
   126  	logger := logs_.GetLogger(ctx)
   127  	logger = logger.WithField("trace_id", span.SpanContext().TraceID()).WithField("span_id", span.SpanContext().SpanID()).WithField("upload_url", uploadUrl)
   128  
   129  	proxy := f.getProxy()
   130  
   131  	opts := []http_.ClientOption{http_.WithDisableKeepAlives(true)}
   132  	if proxy.TargetHost != "" {
   133  		opts = append(opts, http_.WithTargetHost(proxy.TargetHost))
   134  	} else {
   135  
   136  		if proxy.ProxyUrl != "" {
   137  			opts = append(opts, http_.WithProxyURL(proxy.ProxyUrl))
   138  		}
   139  		if proxy.ProxyHost != "" {
   140  			opts = append(opts, http_.WithProxyHost(proxy.ProxyHost))
   141  		}
   142  	}
   143  	opts = append(opts, http_.WithTimeout(f.opts.uploadTimeout))
   144  
   145  	err = time_.RetryWithContext(ctx, func(ctx context.Context) error {
   146  		client, err := http_.NewClient(opts...)
   147  		if err != nil {
   148  			logger.WithError(err).Errorf("new http client err: %v", err)
   149  			return err
   150  		}
   151  		data, err = client.Put(ctx, uploadUrl, "", nil, body)
   152  		if err != nil {
   153  			logger.WithError(err).Errorf("http client put err: %v", err)
   154  			return err
   155  		}
   156  		return nil
   157  
   158  	}, f.opts.retryInterval, f.opts.retryTimes)
   159  
   160  	return data, err
   161  }