github.com/uber/kraken@v0.1.4/lib/torrent/scheduler/dispatch/sync_bitfield.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package dispatch
    15  
    16  import (
    17  	"bytes"
    18  	"sync"
    19  
    20  	"github.com/willf/bitset"
    21  )
    22  
    23  type syncBitfield struct {
    24  	sync.RWMutex
    25  	b *bitset.BitSet
    26  }
    27  
    28  func newSyncBitfield(b *bitset.BitSet) *syncBitfield {
    29  	return &syncBitfield{
    30  		b: b.Clone(),
    31  	}
    32  }
    33  
    34  func (s *syncBitfield) Copy() *bitset.BitSet {
    35  	s.RLock()
    36  	defer s.RUnlock()
    37  
    38  	b := &bitset.BitSet{}
    39  	s.b.Copy(b)
    40  	return b
    41  }
    42  
    43  func (s *syncBitfield) Intersection(other *bitset.BitSet) *bitset.BitSet {
    44  	s.RLock()
    45  	defer s.RUnlock()
    46  
    47  	return s.b.Intersection(other)
    48  }
    49  
    50  func (s *syncBitfield) Len() uint {
    51  	s.RLock()
    52  	defer s.RUnlock()
    53  
    54  	return s.b.Len()
    55  }
    56  
    57  func (s *syncBitfield) Has(i uint) bool {
    58  	s.RLock()
    59  	defer s.RUnlock()
    60  
    61  	return s.b.Test(i)
    62  }
    63  
    64  func (s *syncBitfield) Complete() bool {
    65  	s.RLock()
    66  	defer s.RUnlock()
    67  
    68  	return s.b.All()
    69  }
    70  
    71  func (s *syncBitfield) Set(i uint, v bool) {
    72  	s.Lock()
    73  	defer s.Unlock()
    74  
    75  	s.b.SetTo(i, v)
    76  }
    77  
    78  // GetAllSet returns the indices of all set bits in the bitset.
    79  func (s *syncBitfield) GetAllSet() []uint {
    80  	s.RLock()
    81  	defer s.RUnlock()
    82  
    83  	all := make([]uint, 0, s.b.Len())
    84  	buffer := make([]uint, s.b.Len())
    85  	j := uint(0)
    86  	j, buffer = s.b.NextSetMany(j, buffer)
    87  	for ; len(buffer) > 0; j, buffer = s.b.NextSetMany(j, buffer) {
    88  		all = append(all, buffer...)
    89  		j++
    90  	}
    91  	return all
    92  }
    93  
    94  func (s *syncBitfield) SetAll(v bool) {
    95  	s.Lock()
    96  	defer s.Unlock()
    97  
    98  	for i := uint(0); i < s.b.Len(); i++ {
    99  		s.b.SetTo(i, v)
   100  	}
   101  }
   102  
   103  func (s *syncBitfield) String() string {
   104  	s.RLock()
   105  	defer s.RUnlock()
   106  
   107  	var buf bytes.Buffer
   108  	for i := uint(0); i < s.b.Len(); i++ {
   109  		if s.b.Test(i) {
   110  			buf.WriteString("1")
   111  		} else {
   112  			buf.WriteString("0")
   113  		}
   114  	}
   115  	return buf.String()
   116  }