github.com/0chain/gosdk@v1.17.11/zboxcore/zboxutil/download_buffer.go (about)

     1  package zboxutil
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/0chain/gosdk/core/sys"
     9  	"github.com/valyala/bytebufferpool"
    10  )
    11  
    12  type DownloadBuffer interface {
    13  	RequestChunk(ctx context.Context, num int) []byte
    14  	ReleaseChunk(num int)
    15  	ClearBuffer()
    16  }
    17  
    18  type DownloadBufferWithChan struct {
    19  	buf     []byte
    20  	length  int
    21  	reqSize int
    22  	ch      chan int
    23  	mu      sync.Mutex
    24  	mp      map[int]int
    25  }
    26  
    27  func NewDownloadBufferWithChan(size, numBlocks, effectiveBlockSize int) *DownloadBufferWithChan {
    28  	numBlocks++
    29  	db := &DownloadBufferWithChan{
    30  		buf:     make([]byte, size*numBlocks*effectiveBlockSize),
    31  		length:  size,
    32  		reqSize: effectiveBlockSize * numBlocks,
    33  		ch:      make(chan int, size),
    34  		mp:      make(map[int]int),
    35  	}
    36  	for i := 0; i < size; i++ {
    37  		db.ch <- i
    38  	}
    39  	return db
    40  }
    41  
    42  func (r *DownloadBufferWithChan) ReleaseChunk(num int) {
    43  	r.mu.Lock()
    44  	ind, ok := r.mp[num]
    45  	if !ok {
    46  		r.mu.Unlock()
    47  		return
    48  	}
    49  	delete(r.mp, num)
    50  	r.mu.Unlock()
    51  	r.ch <- ind
    52  }
    53  
    54  func (r *DownloadBufferWithChan) RequestChunk(ctx context.Context, num int) []byte {
    55  	select {
    56  	case <-ctx.Done():
    57  		return nil
    58  	case ind := <-r.ch:
    59  		r.mu.Lock()
    60  		r.mp[num] = ind
    61  		r.mu.Unlock()
    62  		return r.buf[ind*r.reqSize : (ind+1)*r.reqSize : (ind+1)*r.reqSize]
    63  	}
    64  }
    65  
    66  func (r *DownloadBufferWithChan) ClearBuffer() {
    67  	r.buf = nil
    68  	close(r.ch)
    69  	for k := range r.mp {
    70  		delete(r.mp, k)
    71  	}
    72  	r.mp = nil
    73  }
    74  
    75  type DownloadBufferWithMask struct {
    76  	downloadBuf []*bytebufferpool.ByteBuffer
    77  	length      int
    78  	reqSize     int
    79  	numBlocks   int
    80  	mask        uint32
    81  	mu          sync.Mutex
    82  }
    83  
    84  func NewDownloadBufferWithMask(size, numBlocks, effectiveBlockSize int) *DownloadBufferWithMask {
    85  	numBlocks++
    86  	return &DownloadBufferWithMask{
    87  		length:      size,
    88  		reqSize:     effectiveBlockSize * numBlocks,
    89  		mask:        (1 << size) - 1,
    90  		downloadBuf: make([]*bytebufferpool.ByteBuffer, size),
    91  	}
    92  }
    93  
    94  func (r *DownloadBufferWithMask) SetNumBlocks(numBlocks int) {
    95  	r.numBlocks = numBlocks
    96  }
    97  
    98  func (r *DownloadBufferWithMask) RequestChunk(ctx context.Context, num int) []byte {
    99  	num = num / r.numBlocks
   100  	num = num % r.length
   101  	for {
   102  		select {
   103  		case <-ctx.Done():
   104  			return nil
   105  		default:
   106  		}
   107  		r.mu.Lock()
   108  		isSet := r.mask & (1 << num)
   109  		// already assigned
   110  		if isSet == 0 {
   111  			r.mu.Unlock()
   112  			sys.Sleep(200 * time.Millisecond)
   113  			continue
   114  		}
   115  		// assign the chunk by clearing the bit
   116  		r.mask &= ^(1 << num)
   117  		if r.downloadBuf[num] == nil {
   118  			buff := BufferPool.Get()
   119  			if cap(buff.B) < r.reqSize {
   120  				buff.B = make([]byte, r.reqSize)
   121  			}
   122  			r.downloadBuf[num] = buff
   123  		}
   124  		r.mu.Unlock()
   125  		return r.downloadBuf[num].B[:r.reqSize:r.reqSize]
   126  	}
   127  }
   128  
   129  func (r *DownloadBufferWithMask) ReleaseChunk(num int) {
   130  	num = num / r.numBlocks
   131  	num = num % r.length
   132  	r.mu.Lock()
   133  	defer r.mu.Unlock()
   134  	r.mask |= 1 << num
   135  }
   136  
   137  func (r *DownloadBufferWithMask) ClearBuffer() {
   138  	for _, buff := range r.downloadBuf {
   139  		if buff != nil {
   140  			BufferPool.Put(buff)
   141  		}
   142  	}
   143  	r.downloadBuf = nil
   144  }