feat: handle large file uploading while the file chunk is already created

Author: HFO4

.gitignore

@@ -4,6 +4,7 @@
 *.dll
 *.so
 *.dylib
+debug
 
 # Test binary, build with `go test -c`
 *.test

conf.yaml

@@ -1,3 +1,4 @@
 token: 52121225
 api: http://127.0.0.1/Queue
-taskNum: 4
+taskNum: 4
+threadNum: 4

debug

@@ -54,7 +54,9 @@ func main() {
 			}
 			for {
 				taskListContent := api.GetTaskList(config.TASKNUM)
-				task.Init(taskListContent, api, siteInfo)
+				if taskListContent != "none" {
+					task.Init(taskListContent, api, siteInfo)
+				}
 				break
 			}
 		}

main.go

@@ -1,6 +1,7 @@
 package onedrive
 
 import (
+	"bytes"
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
@@ -40,6 +41,13 @@ type Response struct {
 	ResString string
 }
 
+type uploadSessionResponse struct {
+	DataContxt         string   `json:"@odata.context"`
+	ExpirationDateTime string   `json:"expirationDateTime"`
+	NextExpectedRanges []string `json:"nextExpectedRanges"`
+	UploadURL          string   `json:"uploadUrl"`
+}
+
 //Init 初始化客户端
 func (client *Client) Init() bool {
 	client.HTTPClient = &http.Client{}
@@ -55,6 +63,47 @@ func (client *Client) PutFile(path string, file *os.File) (string, string) {
 	return "", res.Error.Error.Message
 }
 
+//CreateUploadSession 创建分片上传会话
+func (client *Client) CreateUploadSession(path string) (string, string) {
+	res := client.apiPost(path, []byte(""))
+	if res.Success {
+		response := uploadSessionResponse{}
+		err := json.Unmarshal([]byte(res.ResString), &response)
+		if err != nil {
+			return "", err.Error()
+		}
+		return response.UploadURL, ""
+	}
+
+	return "", res.Error.Error.Message
+}
+
+//apiPost 发送POST请求
+func (client *Client) apiPost(path string, jsonStr []byte) Response {
+	if client.Tried > maxTry {
+		return buildResponseResult("PUT failed, reached the maximum number of attempts.", 0)
+	}
+
+	req, err := http.NewRequest("POST", APIURL+path, bytes.NewBuffer(jsonStr))
+	if err != nil {
+		return buildResponseResult(err.Error(), 0)
+	}
+
+	req.Header.Set("Authorization", "Bearer "+client.AccessToken)
+	req.Header.Set("Content-Type", "application/json")
+
+	res, err := client.HTTPClient.Do(req)
+	if err != nil {
+		client.Tried++
+		return client.apiPost(path, jsonStr)
+	}
+	defer res.Body.Close()
+
+	r, _ := ioutil.ReadAll(res.Body)
+	return client.praseResponse(string(r), res.StatusCode)
+
+}
+
 //apiPut 发送PUT请求
 func (client *Client) apiPut(path string, stream *os.File) Response {
 	if client.Tried > maxTry {
@@ -80,7 +129,7 @@ func (client *Client) apiPut(path string, stream *os.File) Response {
 }
 
 func (client *Client) praseResponse(res string, code int) Response {
-	if code != 200 {
+	if code != 200 && code != 202 {
 		errorRes := ErrorResponse{}
 		json.Unmarshal([]byte(res), &errorRes)
 		return Response{

onedrive/client.go

@@ -2,21 +2,36 @@ package task
 
 import (
 	"encoding/json"
+	"fmt"
 	"log"
 	"os"
 	"strconv"
+	"sync"
 
 	"../onedrive"
 )
 
+const uploadThreadNum int = 4
+
+type chunkFile struct {
+	ID      int    `json:"id"`
+	User    int    `json:"user"`
+	CTX     string `json:"ctx"`
+	Time    string `json:"time"`
+	ObjName string `json:"obj_name"`
+	ChunkID int    `json:"chunk_id"`
+	Sum     int    `json:"sum"`
+}
+
 type oneDriveUploadAttr struct {
-	Fname    string `json:"fname"`
-	Path     string `json:"path"`
-	Objname  string `json:"objname"`
-	SavePath string `json:"savePath"`
-	Fsize    uint64 `json:"fsize"`
-	PicInfo  string `json:"picInfo"`
-	PolicyID int    `json:"policyId"`
+	Fname    string      `json:"fname"`
+	Path     string      `json:"path"`
+	Objname  string      `json:"objname"`
+	SavePath string      `json:"savePath"`
+	Fsize    uint64      `json:"fsize"`
+	PicInfo  string      `json:"picInfo"`
+	PolicyID int         `json:"policyId"`
+	Chunks   []chunkFile `json:"chunks"`
 }
 
 //OneDriveUpload OneDrive上传类型Task
@@ -44,6 +59,14 @@ type OnedriveState struct {
 	} `json:"token"`
 }
 
+//Chunk 文件分片
+type Chunk struct {
+	Type      int //分片方式，0-已切割为文件片，1-无需切割
+	From      int
+	To        int
+	ChunkPath string
+}
+
 //Excute 执行Onedrive上传Task
 func (task *OneDriveUpload) Excute() {
 	authState := OnedriveState{}
@@ -55,13 +78,103 @@ func (task *OneDriveUpload) Excute() {
 		Tried:        0,
 	}
 	Client.Init()
-	var filePath string
-	if task.Type == "UploadRegularRemoteDownloadFileToOnedrive" {
 
+	if task.Type == "uploadSingleToOnedrive" {
+		task.uploadRegularFile(&Client)
+	} else if task.Type == "uploadChunksToOnedrive" {
+		task.uploadChunks(&Client)
+	}
+
+}
+
+func (task *OneDriveUpload) uploadChunks(Client *onedrive.Client) {
+	//获取上传URL
+	url, err := Client.CreateUploadSession("/me/drive/root:/" + task.Attr.SavePath + "/" + task.Attr.Objname + ":/createUploadSession")
+	if err != "" {
+		task.Log("[Error] Failed to create upload session," + err)
+		task.Error()
+		return
+	}
+	fmt.Println(url)
+
+	chunkList, chunkErr := task.buildChunks()
+	if chunkErr != "" {
+		task.Log("[Error] Failed to upload chunks," + chunkErr)
+		task.Error()
+		return
+	}
+
+	var wg sync.WaitGroup
+	ch := make(chan Chunk)
+	isFailed := false
+
+	for index := 0; index < uploadThreadNum; index++ {
+		wg.Add(1)
+		go task.uploadSingleChunk(&wg, ch, &isFailed)
+	}
+
+	for _, v := range chunkList {
+		if isFailed {
+			close(ch)
+			break
+		}
+		ch <- v
+	}
+	close(ch)
+	wg.Wait()
+
+}
+
+func (task *OneDriveUpload) uploadSingleChunk(wg *sync.WaitGroup, ch chan Chunk, isFailed *bool) {
+	for {
+		chunk, opened := <-ch
+		if !opened {
+			fmt.Println("quit")
+			wg.Done()
+			return
+		}
+		fmt.Println(chunk.From)
+	}
+}
+
+func (task *OneDriveUpload) buildChunks() ([]Chunk, string) {
+
+	var chunkType int
+	if task.Type == "uploadChunksToOnedrive" {
+		chunkType = 0
 	} else {
-		filePath = task.BasePath + "public/uploads/" + task.Attr.SavePath + "/" + task.Attr.Objname
+		chunkType = 1
 	}
 
+	var chunkList []Chunk
+	var offset int
+
+	for _, v := range task.Attr.Chunks {
+
+		chunkPath := task.BasePath + "public/uploads/chunks/" + v.ObjName + ".chunk"
+
+		fileInfo, err := os.Stat(chunkPath)
+		if os.IsNotExist(err) {
+			return chunkList, "Chunk file " + chunkPath + " not exist"
+		}
+		chunkSize := int(fileInfo.Size())
+
+		chunkList = append(chunkList, Chunk{
+			Type:      chunkType,
+			From:      offset,
+			To:        offset + chunkSize - 1,
+			ChunkPath: chunkPath,
+		})
+		offset += chunkSize
+	}
+
+	return chunkList, ""
+
+}
+
+func (task *OneDriveUpload) uploadRegularFile(Client *onedrive.Client) {
+	var filePath string
+	filePath = task.BasePath + "public/uploads/" + task.Attr.SavePath + "/" + task.Attr.Objname
 	r, err := os.Open(filePath)
 	defer r.Close()
 	if err != nil {
@@ -74,14 +187,14 @@ func (task *OneDriveUpload) Excute() {
 	if errorMsg != "" {
 		task.Log("[Error] Upload Failed," + errorMsg)
 		task.Error()
+		return
 	}
 
 	addRes := task.Info.apiInfo.SetSuccess(task.Info.sqlInfo.ID)
 
 	if addRes != "" {
 		task.Log("[Error] " + addRes)
 	}
-
 }
 
 //Init 执行Onedrive上传Task

task/TaskAdapter.go

@@ -52,6 +52,14 @@ func (task *taskInfo) Run(wg *sync.WaitGroup) {
 			BasePath: task.siteInfo["basePath"],
 		}
 		newTask.Init()
+	case task.sqlInfo.TaskType == "uploadChunksToOnedrive":
+		newTask = &OneDriveUpload{
+			Info:     task,
+			Tried:    0,
+			Type:     task.sqlInfo.TaskType,
+			BasePath: task.siteInfo["basePath"],
+		}
+		newTask.Init()
 	}
 	newTask.Excute()
 	wg.Done()