github.com/0chain/gosdk@v1.17.11/zboxcore/sdk/download_progress_storer.go (about)

     1  package sdk
     2  
     3  import (
     4  	"container/heap"
     5  	"context"
     6  	"encoding/json"
     7  	"os"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/0chain/errors"
    12  	"github.com/0chain/gosdk/core/sys"
    13  )
    14  
    15  type DownloadProgressStorer interface {
    16  	// Load load download progress by id
    17  	Load(id string, numBlocks int) *DownloadProgress
    18  	// Update download progress
    19  	Update(writtenBlock int)
    20  	// Remove remove download progress by id
    21  	Remove() error
    22  	// Start start download progress
    23  	Start(ctx context.Context)
    24  	// Save download progress
    25  	Save(dp *DownloadProgress)
    26  }
    27  
    28  type FsDownloadProgressStorer struct {
    29  	sync.Mutex
    30  	isRemoved bool
    31  	dp        *DownloadProgress
    32  	next      int
    33  	queue     queue
    34  }
    35  
    36  // CreateFsDownloadProgress create a download progress storer instance to track download progress and queue
    37  func CreateFsDownloadProgress() *FsDownloadProgressStorer {
    38  	down := &FsDownloadProgressStorer{
    39  		queue: make(queue, 0),
    40  	}
    41  	heap.Init(&down.queue)
    42  	return down
    43  }
    44  
    45  func (ds *FsDownloadProgressStorer) Start(ctx context.Context) {
    46  	tc := time.NewTicker(2 * time.Second)
    47  	ds.next += ds.dp.numBlocks
    48  	go func() {
    49  		defer tc.Stop()
    50  		for {
    51  			select {
    52  			case <-ctx.Done():
    53  				return
    54  			case <-tc.C:
    55  				ds.Lock()
    56  				if ds.isRemoved {
    57  					ds.Unlock()
    58  					return
    59  				}
    60  				if len(ds.queue) > 0 && ds.queue[0] == ds.next {
    61  					for len(ds.queue) > 0 && ds.queue[0] == ds.next {
    62  						ds.dp.LastWrittenBlock = ds.next
    63  						heap.Pop(&ds.queue)
    64  						ds.next += ds.dp.numBlocks
    65  					}
    66  					ds.Unlock()
    67  					ds.saveToDisk()
    68  				} else {
    69  					ds.Unlock()
    70  				}
    71  			}
    72  		}
    73  	}()
    74  }
    75  
    76  func (ds *FsDownloadProgressStorer) Load(progressID string, numBlocks int) *DownloadProgress {
    77  	dp := &DownloadProgress{}
    78  	buf, err := sys.Files.LoadProgress(progressID)
    79  	if err != nil {
    80  		return nil
    81  	}
    82  	if err = json.Unmarshal(buf, dp); err != nil {
    83  		return nil
    84  	}
    85  	ds.dp = dp
    86  	dp.numBlocks = numBlocks
    87  	ds.next = dp.LastWrittenBlock
    88  	return ds.dp
    89  }
    90  
    91  func (ds *FsDownloadProgressStorer) saveToDisk() {
    92  	ds.Lock()
    93  	defer ds.Unlock()
    94  	if ds.isRemoved {
    95  		return
    96  	}
    97  	buf, err := json.Marshal(ds.dp)
    98  	if err != nil {
    99  		return
   100  	}
   101  	err = sys.Files.SaveProgress(ds.dp.ID, buf, 0666)
   102  	if err != nil {
   103  		return
   104  	}
   105  }
   106  
   107  func (ds *FsDownloadProgressStorer) Save(dp *DownloadProgress) {
   108  	ds.dp = dp
   109  	ds.saveToDisk()
   110  }
   111  
   112  func (ds *FsDownloadProgressStorer) Update(writtenBlock int) {
   113  	ds.Lock()
   114  	defer ds.Unlock()
   115  	if ds.isRemoved {
   116  		return
   117  	}
   118  	heap.Push(&ds.queue, writtenBlock)
   119  }
   120  
   121  func (ds *FsDownloadProgressStorer) Remove() error {
   122  	ds.Lock()
   123  	defer ds.Unlock()
   124  	if ds.isRemoved || ds.dp == nil {
   125  		return nil
   126  	}
   127  	ds.isRemoved = true
   128  	err := sys.Files.RemoveProgress(ds.dp.ID)
   129  	if err != nil && !errors.Is(err, os.ErrNotExist) {
   130  		return err
   131  	}
   132  	return nil
   133  }