github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/internal/ackseq/ackseq.go (about)

     1  // Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package ackseq
     6  
     7  import (
     8  	"sync"
     9  	"sync/atomic"
    10  
    11  	"github.com/cockroachdb/errors"
    12  )
    13  
    14  const (
    15  	// The window size constants. These values specify a window that can hold ~1m
    16  	// pending unacknowledged sequence numbers using 128 KB of memory.
    17  	windowSize  = 1 << 20
    18  	windowMask  = windowSize - 1
    19  	windowBytes = (windowSize + 7) / 8
    20  )
    21  
    22  // S keeps track of the largest sequence number such that all sequence numbers
    23  // in the range [0,v) have been acknowledged.
    24  type S struct {
    25  	next uint64
    26  	mu   struct {
    27  		sync.Mutex
    28  		base   uint64
    29  		window [windowBytes]uint8
    30  	}
    31  }
    32  
    33  // New creates a new acknowledged sequence tracker with the specified base
    34  // sequence number. All of the sequence numbers in the range [0,base) are
    35  // considered acknowledged. Next() will return base upon first call.
    36  func New(base uint64) *S {
    37  	s := &S{
    38  		next: base,
    39  	}
    40  	s.mu.base = base
    41  	return s
    42  }
    43  
    44  // Next returns the next sequence number to use.
    45  func (s *S) Next() uint64 {
    46  	return atomic.AddUint64(&s.next, 1) - 1
    47  }
    48  
    49  // Ack acknowledges the specified seqNum, adjusting base as necessary,
    50  // returning the number of newly acknowledged sequence numbers.
    51  func (s *S) Ack(seqNum uint64) (int, error) {
    52  	s.mu.Lock()
    53  	if s.getLocked(seqNum) {
    54  		defer s.mu.Unlock()
    55  		return 0, errors.Errorf(
    56  			"pending acks exceeds window size: %d has been acked, but %d has not",
    57  			errors.Safe(seqNum), errors.Safe(s.mu.base))
    58  	}
    59  
    60  	var count int
    61  	s.setLocked(seqNum)
    62  	for s.getLocked(s.mu.base) {
    63  		s.clearLocked(s.mu.base)
    64  		s.mu.base++
    65  		count++
    66  	}
    67  	s.mu.Unlock()
    68  	return count, nil
    69  }
    70  
    71  func (s *S) getLocked(seqNum uint64) bool {
    72  	bit := seqNum & windowMask
    73  	return (s.mu.window[bit/8] & (1 << (bit % 8))) != 0
    74  }
    75  
    76  func (s *S) setLocked(seqNum uint64) {
    77  	bit := seqNum & windowMask
    78  	s.mu.window[bit/8] |= (1 << (bit % 8))
    79  }
    80  
    81  func (s *S) clearLocked(seqNum uint64) {
    82  	bit := seqNum & windowMask
    83  	s.mu.window[bit/8] &^= (1 << (bit % 8))
    84  }