github.com/qjfoidnh/BaiduPCS-Go@v0.0.0-20231011165705-caa18a3765f3/requester/downloader/instance_state.go (about)

     1  package downloader
     2  
     3  import (
     4  	"errors"
     5  	"github.com/golang/protobuf/proto"
     6  	"github.com/qjfoidnh/BaiduPCS-Go/pcsutil/cachepool"
     7  	"github.com/qjfoidnh/BaiduPCS-Go/pcsverbose"
     8  	"github.com/qjfoidnh/BaiduPCS-Go/requester/transfer"
     9  	"github.com/json-iterator/go"
    10  	"os"
    11  	"sync"
    12  )
    13  
    14  type (
    15  	//InstanceState 状态, 断点续传信息
    16  	InstanceState struct {
    17  		saveFile *os.File
    18  		format   InstanceStateStorageFormat
    19  		ii       transfer.DownloadInstanceInfoExporter
    20  		mu       sync.Mutex
    21  	}
    22  
    23  	// InstanceStateStorageFormat 断点续传储存类型
    24  	InstanceStateStorageFormat int
    25  )
    26  
    27  const (
    28  	// InstanceStateStorageFormatJSON json 格式
    29  	InstanceStateStorageFormatJSON = iota
    30  	// InstanceStateStorageFormatProto3 protobuf 格式
    31  	InstanceStateStorageFormatProto3
    32  )
    33  
    34  //NewInstanceState 初始化InstanceState
    35  func NewInstanceState(saveFile *os.File, format InstanceStateStorageFormat) *InstanceState {
    36  	return &InstanceState{
    37  		saveFile: saveFile,
    38  		format:   format,
    39  	}
    40  }
    41  
    42  func (is *InstanceState) checkSaveFile() bool {
    43  	return is.saveFile != nil
    44  }
    45  
    46  func (is *InstanceState) getSaveFileContents() []byte {
    47  	if !is.checkSaveFile() {
    48  		return nil
    49  	}
    50  
    51  	finfo, err := is.saveFile.Stat()
    52  	if err != nil {
    53  		panic(err)
    54  	}
    55  
    56  	size := finfo.Size()
    57  	if size > 0xffffffff {
    58  		panic("savePath too large")
    59  	}
    60  	intSize := int(size)
    61  
    62  	buf := cachepool.RawMallocByteSlice(intSize)
    63  
    64  	n, _ := is.saveFile.ReadAt(buf, 0)
    65  	return buf[:n]
    66  }
    67  
    68  //Get 获取断点续传信息
    69  func (is *InstanceState) Get() (eii *transfer.DownloadInstanceInfo) {
    70  	if !is.checkSaveFile() {
    71  		return nil
    72  	}
    73  
    74  	is.mu.Lock()
    75  	defer is.mu.Unlock()
    76  
    77  	contents := is.getSaveFileContents()
    78  	if len(contents) <= 0 {
    79  		return
    80  	}
    81  
    82  	is.ii = &transfer.DownloadInstanceInfoExport{}
    83  	var err error
    84  	switch is.format {
    85  	case InstanceStateStorageFormatProto3:
    86  		err = proto.Unmarshal(contents, is.ii.(*transfer.DownloadInstanceInfoExport))
    87  	default:
    88  		err = jsoniter.Unmarshal(contents, is.ii)
    89  	}
    90  
    91  	if err != nil {
    92  		pcsverbose.Verbosef("DEBUG: InstanceInfo unmarshal error: %s\n", err)
    93  		return
    94  	}
    95  
    96  	eii = is.ii.GetInstanceInfo()
    97  	return
    98  }
    99  
   100  //Put 提交断点续传信息
   101  func (is *InstanceState) Put(eii *transfer.DownloadInstanceInfo) {
   102  	if !is.checkSaveFile() {
   103  		return
   104  	}
   105  
   106  	is.mu.Lock()
   107  	defer is.mu.Unlock()
   108  
   109  	if is.ii == nil {
   110  		is.ii = &transfer.DownloadInstanceInfoExport{}
   111  	}
   112  	is.ii.SetInstanceInfo(eii)
   113  	var (
   114  		data []byte
   115  		err  error
   116  	)
   117  	switch is.format {
   118  	case InstanceStateStorageFormatProto3:
   119  		data, err = proto.Marshal(is.ii.(*transfer.DownloadInstanceInfoExport))
   120  	default:
   121  		data, err = jsoniter.Marshal(is.ii)
   122  	}
   123  	if err != nil {
   124  		panic(err)
   125  	}
   126  
   127  	err = is.saveFile.Truncate(int64(len(data)))
   128  	if err != nil {
   129  		pcsverbose.Verbosef("DEBUG: truncate file error: %s\n", err)
   130  	}
   131  
   132  	_, err = is.saveFile.WriteAt(data, 0)
   133  	if err != nil {
   134  		pcsverbose.Verbosef("DEBUG: write instance state error: %s\n", err)
   135  	}
   136  }
   137  
   138  //Close 关闭
   139  func (is *InstanceState) Close() error {
   140  	if !is.checkSaveFile() {
   141  		return nil
   142  	}
   143  
   144  	return is.saveFile.Close()
   145  }
   146  
   147  func (der *Downloader) initInstanceState(format InstanceStateStorageFormat) (err error) {
   148  	if der.instanceState != nil {
   149  		return errors.New("already initInstanceState")
   150  	}
   151  
   152  	var saveFile *os.File
   153  	if !der.config.IsTest && der.config.InstanceStatePath != "" {
   154  		saveFile, err = os.OpenFile(der.config.InstanceStatePath, os.O_RDWR|os.O_CREATE, 0777)
   155  		if err != nil {
   156  			return err
   157  		}
   158  	}
   159  
   160  	der.instanceState = NewInstanceState(saveFile, format)
   161  	return nil
   162  }
   163  
   164  func (der *Downloader) removeInstanceState() error {
   165  	der.instanceState.Close()
   166  	if !der.config.IsTest && der.config.InstanceStatePath != "" {
   167  		return os.Remove(der.config.InstanceStatePath)
   168  	}
   169  	return nil
   170  }