github.com/anacrolix/torrent@v1.61.0/storage/map-piece-completion.go (about)

     1  package storage
     2  
     3  import (
     4  	"iter"
     5  	"sync"
     6  
     7  	g "github.com/anacrolix/generics"
     8  
     9  	"github.com/anacrolix/torrent/metainfo"
    10  	"github.com/anacrolix/torrent/types/infohash"
    11  )
    12  
    13  type mapPieceCompletion struct {
    14  	// TODO: Generics. map of InfoHash to *memoryTorrentJustComplete.
    15  	m sync.Map
    16  }
    17  
    18  func (me *mapPieceCompletion) Persistent() bool {
    19  	return false
    20  }
    21  
    22  type (
    23  	justComplete              = g.Option[bool]
    24  	memoryTorrentJustComplete struct {
    25  		mu    sync.RWMutex
    26  		state []justComplete
    27  	}
    28  )
    29  
    30  func (me *memoryTorrentJustComplete) Get(i int) justComplete {
    31  	me.mu.RLock()
    32  	defer me.mu.RUnlock()
    33  	return me.getLocked(i)
    34  }
    35  
    36  func (me *memoryTorrentJustComplete) getLocked(i int) justComplete {
    37  	if i >= len(me.state) {
    38  		return g.None[bool]()
    39  	}
    40  	return me.state[i]
    41  }
    42  
    43  func (me *memoryTorrentJustComplete) Set(i int, complete bool) {
    44  	me.mu.Lock()
    45  	defer me.mu.Unlock()
    46  	for i >= len(me.state) {
    47  		me.state = append(me.state, g.None[bool]())
    48  	}
    49  	me.state[i].Set(complete)
    50  }
    51  
    52  func (me *memoryTorrentJustComplete) GetRange(begin, end int) iter.Seq[justComplete] {
    53  	me.mu.RLock()
    54  	return func(yield func(justComplete) bool) {
    55  		defer me.mu.RUnlock()
    56  		for i := begin; i < end; i++ {
    57  			if !yield(me.getLocked(i)) {
    58  				return
    59  			}
    60  		}
    61  	}
    62  }
    63  
    64  var _ interface {
    65  	PieceCompletion
    66  	PieceCompletionGetRanger
    67  } = (*mapPieceCompletion)(nil)
    68  
    69  func NewMapPieceCompletion() PieceCompletion {
    70  	return &mapPieceCompletion{}
    71  }
    72  
    73  func (me *mapPieceCompletion) Close() error {
    74  	me.m.Clear()
    75  	return nil
    76  }
    77  
    78  func (me *mapPieceCompletion) Get(pk metainfo.PieceKey) (c Completion, err error) {
    79  	v, ok := me.m.Load(pk.InfoHash)
    80  	if !ok {
    81  		return
    82  	}
    83  	jcs := v.(*memoryTorrentJustComplete)
    84  	c.Complete, c.Ok = jcs.Get(pk.Index).AsTuple()
    85  	return
    86  }
    87  
    88  func (me *mapPieceCompletion) Set(pk metainfo.PieceKey, complete bool) error {
    89  	v, ok := me.m.Load(pk.InfoHash)
    90  	if !ok {
    91  		v, _ = me.m.LoadOrStore(pk.InfoHash, &memoryTorrentJustComplete{})
    92  	}
    93  	t := v.(*memoryTorrentJustComplete)
    94  	t.Set(pk.Index, complete)
    95  	return nil
    96  }
    97  
    98  func (me *mapPieceCompletion) GetRange(ih infohash.T, begin, end int) iter.Seq[Completion] {
    99  	return func(yield func(Completion) bool) {
   100  		v, ok := me.m.Load(ih)
   101  		if !ok {
   102  			return
   103  		}
   104  		t := v.(*memoryTorrentJustComplete)
   105  		for jc := range t.GetRange(begin, end) {
   106  			if !yield(Completion{
   107  				Err:      nil,
   108  				Ok:       jc.Ok,
   109  				Complete: jc.Value,
   110  			}) {
   111  				return
   112  			}
   113  		}
   114  	}
   115  }