github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/topicsdb/sync.go (about)

     1  package topicsdb
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  )
     7  
     8  type (
     9  	posCounter struct {
    10  		pos   int
    11  		count int
    12  	}
    13  
    14  	blockCounter struct {
    15  		wait  chan struct{}
    16  		count int
    17  	}
    18  
    19  	synchronizator struct {
    20  		mu        sync.Mutex
    21  		threads   sync.WaitGroup
    22  		positions map[int]*posCounter
    23  		goNext    bool
    24  		minBlock  uint64
    25  		blocks    map[uint64]*blockCounter
    26  	}
    27  )
    28  
    29  func newSynchronizator() *synchronizator {
    30  	s := &synchronizator{
    31  		positions: make(map[int]*posCounter),
    32  		goNext:    true,
    33  		minBlock:  0,
    34  		blocks:    make(map[uint64]*blockCounter),
    35  	}
    36  
    37  	return s
    38  }
    39  
    40  func (s *synchronizator) Halt() {
    41  	s.goNext = false
    42  
    43  	s.mu.Lock()
    44  	defer s.mu.Unlock()
    45  
    46  	for n := range s.blocks {
    47  		if n != s.minBlock {
    48  			close(s.blocks[n].wait)
    49  		}
    50  	}
    51  }
    52  
    53  func (s *synchronizator) GoNext(n uint64) (prev uint64, gonext bool) {
    54  	if !s.goNext {
    55  		return
    56  	}
    57  
    58  	if n > s.minBlock {
    59  		s.mu.Lock()
    60  		prev = s.minBlock
    61  		s.enqueueBlock(n)
    62  		s.dequeueBlock()
    63  		wait := s.blocks[n].wait
    64  		s.mu.Unlock()
    65  		// wait for other threads
    66  		<-wait
    67  	}
    68  
    69  	gonext = s.goNext
    70  	return
    71  }
    72  
    73  func (s *synchronizator) StartThread(pos int, num int) {
    74  	s.threads.Add(1)
    75  
    76  	s.mu.Lock()
    77  	defer s.mu.Unlock()
    78  
    79  	s.enqueueBlock(s.minBlock)
    80  
    81  	if _, ok := s.positions[pos]; ok {
    82  		s.positions[pos].count++
    83  	} else {
    84  		s.positions[pos] = &posCounter{pos, 1}
    85  	}
    86  }
    87  
    88  func (s *synchronizator) FinishThread(pos int, num int) {
    89  	s.threads.Done()
    90  
    91  	s.mu.Lock()
    92  	defer s.mu.Unlock()
    93  
    94  	s.positions[pos].count--
    95  	s.dequeueBlock()
    96  }
    97  
    98  func (s *synchronizator) enqueueBlock(n uint64) {
    99  	if _, ok := s.blocks[n]; ok {
   100  		s.blocks[n].count++
   101  	} else {
   102  		s.blocks[n] = &blockCounter{
   103  			wait:  make(chan struct{}),
   104  			count: 1,
   105  		}
   106  	}
   107  }
   108  
   109  func (s *synchronizator) dequeueBlock() {
   110  	s.blocks[s.minBlock].count--
   111  	if s.blocks[s.minBlock].count < 1 {
   112  
   113  		for _, pos := range s.positions {
   114  			if pos.count < 1 {
   115  				s.goNext = false
   116  				break
   117  			}
   118  		}
   119  
   120  		delete(s.blocks, s.minBlock)
   121  		if len(s.blocks) < 1 {
   122  			return
   123  		}
   124  		// find new minBlock
   125  		s.minBlock = math.MaxUint64
   126  		for b := range s.blocks {
   127  			if s.minBlock > b {
   128  				s.minBlock = b
   129  			}
   130  		}
   131  		close(s.blocks[s.minBlock].wait)
   132  	}
   133  }
   134  
   135  func (s *synchronizator) WaitForThreads() {
   136  	s.threads.Wait()
   137  }
   138  
   139  func (s *synchronizator) PositionsCount() int {
   140  	return len(s.positions)
   141  }