github.com/fzfile/BaiduPCS-Go@v0.0.0-20200606205115-4408961cf336/requester/uploader/uploader.go (about)

     1  // Package uploader 上传包
     2  package uploader
     3  
     4  import (
     5  	"github.com/fzfile/BaiduPCS-Go/pcsutil"
     6  	"github.com/fzfile/BaiduPCS-Go/pcsutil/converter"
     7  	"github.com/fzfile/BaiduPCS-Go/pcsverbose"
     8  	"github.com/fzfile/BaiduPCS-Go/requester"
     9  	"github.com/fzfile/BaiduPCS-Go/requester/rio"
    10  	"net/http"
    11  	"time"
    12  )
    13  
    14  const (
    15  	// BufioReadSize bufio 缓冲区大小, 用于上传时读取文件
    16  	BufioReadSize = int(64 * converter.KB) // 64KB
    17  )
    18  
    19  type (
    20  	//CheckFunc 上传完成的检测函数
    21  	CheckFunc func(resp *http.Response, uploadErr error)
    22  
    23  	// Uploader 上传
    24  	Uploader struct {
    25  		url         string   // 上传地址
    26  		readed64    Readed64 // 要上传的对象
    27  		contentType string
    28  
    29  		client *requester.HTTPClient
    30  
    31  		executeTime time.Time
    32  		executed    bool
    33  		finished    chan struct{}
    34  
    35  		checkFunc CheckFunc
    36  		onExecute func()
    37  		onFinish  func()
    38  	}
    39  )
    40  
    41  var (
    42  	uploaderVerbose = pcsverbose.New("UPLOADER")
    43  )
    44  
    45  // NewUploader 返回 uploader 对象, url: 上传地址, readerlen64: 实现 rio.ReaderLen64 接口的对象, 例如文件
    46  func NewUploader(url string, readerlen64 rio.ReaderLen64) (uploader *Uploader) {
    47  	uploader = &Uploader{
    48  		url:      url,
    49  		readed64: NewReaded64(readerlen64),
    50  	}
    51  
    52  	return
    53  }
    54  
    55  func (u *Uploader) lazyInit() {
    56  	if u.finished == nil {
    57  		u.finished = make(chan struct{})
    58  	}
    59  	if u.client == nil {
    60  		u.client = requester.NewHTTPClient()
    61  	}
    62  	u.client.SetTimeout(0)
    63  	u.client.SetResponseHeaderTimeout(0)
    64  }
    65  
    66  // SetClient 设置http客户端
    67  func (u *Uploader) SetClient(c *requester.HTTPClient) {
    68  	u.client = c
    69  }
    70  
    71  //SetContentType 设置Content-Type
    72  func (u *Uploader) SetContentType(contentType string) {
    73  	u.contentType = contentType
    74  }
    75  
    76  //SetCheckFunc 设置上传完成的检测函数
    77  func (u *Uploader) SetCheckFunc(checkFunc CheckFunc) {
    78  	u.checkFunc = checkFunc
    79  }
    80  
    81  // Execute 执行上传, 收到返回值信号则为上传结束
    82  func (u *Uploader) Execute() {
    83  	pcsutil.Trigger(u.onExecute)
    84  
    85  	// 开始上传
    86  	u.executeTime = time.Now()
    87  	u.executed = true
    88  	resp, _, err := u.execute()
    89  
    90  	// 上传结束
    91  	close(u.finished)
    92  
    93  	if u.checkFunc != nil {
    94  		u.checkFunc(resp, err)
    95  	}
    96  
    97  	pcsutil.Trigger(u.onFinish) // 触发上传结束的事件
    98  }
    99  
   100  func (u *Uploader) execute() (resp *http.Response, code int, err error) {
   101  	u.lazyInit()
   102  	header := map[string]string{}
   103  	if u.contentType != "" {
   104  		header["Content-Type"] = u.contentType
   105  	}
   106  
   107  	resp, err = u.client.Req(http.MethodPost, u.url, u.readed64, header)
   108  	if err != nil {
   109  		return nil, 2, err
   110  	}
   111  
   112  	return resp, 0, nil
   113  }
   114  
   115  // OnExecute 任务开始时触发的事件
   116  func (u *Uploader) OnExecute(fn func()) {
   117  	u.onExecute = fn
   118  }
   119  
   120  // OnFinish 任务完成时触发的事件
   121  func (u *Uploader) OnFinish(fn func()) {
   122  	u.onFinish = fn
   123  }