github.com/bytedance/gopkg@v0.0.0-20240514070511-01b2cbcf35e1/collection/skipset/skipset.go (about)

     1  // Copyright 2021 ByteDance 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  
    15  // Package skipset is a high-performance, scalable, concurrent-safe set based on skip-list.
    16  // In the typical pattern(100000 operations, 90%CONTAINS 9%Add 1%Remove, 8C16T), the skipset
    17  // up to 15x faster than the built-in sync.Map.
    18  package skipset
    19  
    20  import (
    21  	"sync"
    22  	"sync/atomic"
    23  	"unsafe"
    24  )
    25  
    26  // Int64Set represents a set based on skip list in ascending order.
    27  type Int64Set struct {
    28  	header       *int64Node
    29  	length       int64
    30  	highestLevel int64 // highest level for now
    31  }
    32  
    33  type int64Node struct {
    34  	value int64
    35  	next  optionalArray // [level]*int64Node
    36  	mu    sync.Mutex
    37  	flags bitflag
    38  	level uint32
    39  }
    40  
    41  func newInt64Node(value int64, level int) *int64Node {
    42  	node := &int64Node{
    43  		value: value,
    44  		level: uint32(level),
    45  	}
    46  	if level > op1 {
    47  		node.next.extra = new([op2]unsafe.Pointer)
    48  	}
    49  	return node
    50  }
    51  
    52  func (n *int64Node) loadNext(i int) *int64Node {
    53  	return (*int64Node)(n.next.load(i))
    54  }
    55  
    56  func (n *int64Node) storeNext(i int, node *int64Node) {
    57  	n.next.store(i, unsafe.Pointer(node))
    58  }
    59  
    60  func (n *int64Node) atomicLoadNext(i int) *int64Node {
    61  	return (*int64Node)(n.next.atomicLoad(i))
    62  }
    63  
    64  func (n *int64Node) atomicStoreNext(i int, node *int64Node) {
    65  	n.next.atomicStore(i, unsafe.Pointer(node))
    66  }
    67  
    68  func (n *int64Node) lessthan(value int64) bool {
    69  	return n.value < value
    70  }
    71  
    72  func (n *int64Node) equal(value int64) bool {
    73  	return n.value == value
    74  }
    75  
    76  // NewInt64 return an empty int64 skip set in ascending order.
    77  func NewInt64() *Int64Set {
    78  	h := newInt64Node(0, maxLevel)
    79  	h.flags.SetTrue(fullyLinked)
    80  	return &Int64Set{
    81  		header:       h,
    82  		highestLevel: defaultHighestLevel,
    83  	}
    84  }
    85  
    86  // findNodeRemove takes a value and two maximal-height arrays then searches exactly as in a sequential skip-list.
    87  // The returned preds and succs always satisfy preds[i] > value >= succs[i].
    88  func (s *Int64Set) findNodeRemove(value int64, preds *[maxLevel]*int64Node, succs *[maxLevel]*int64Node) int {
    89  	// lFound represents the index of the first layer at which it found a node.
    90  	lFound, x := -1, s.header
    91  	for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- {
    92  		succ := x.atomicLoadNext(i)
    93  		for succ != nil && succ.lessthan(value) {
    94  			x = succ
    95  			succ = x.atomicLoadNext(i)
    96  		}
    97  		preds[i] = x
    98  		succs[i] = succ
    99  
   100  		// Check if the value already in the skip list.
   101  		if lFound == -1 && succ != nil && succ.equal(value) {
   102  			lFound = i
   103  		}
   104  	}
   105  	return lFound
   106  }
   107  
   108  // findNodeAdd takes a value and two maximal-height arrays then searches exactly as in a sequential skip-set.
   109  // The returned preds and succs always satisfy preds[i] > value >= succs[i].
   110  func (s *Int64Set) findNodeAdd(value int64, preds *[maxLevel]*int64Node, succs *[maxLevel]*int64Node) int {
   111  	x := s.header
   112  	for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- {
   113  		succ := x.atomicLoadNext(i)
   114  		for succ != nil && succ.lessthan(value) {
   115  			x = succ
   116  			succ = x.atomicLoadNext(i)
   117  		}
   118  		preds[i] = x
   119  		succs[i] = succ
   120  
   121  		// Check if the value already in the skip list.
   122  		if succ != nil && succ.equal(value) {
   123  			return i
   124  		}
   125  	}
   126  	return -1
   127  }
   128  
   129  func unlockInt64(preds [maxLevel]*int64Node, highestLevel int) {
   130  	var prevPred *int64Node
   131  	for i := highestLevel; i >= 0; i-- {
   132  		if preds[i] != prevPred { // the node could be unlocked by previous loop
   133  			preds[i].mu.Unlock()
   134  			prevPred = preds[i]
   135  		}
   136  	}
   137  }
   138  
   139  // Add add the value into skip set, return true if this process insert the value into skip set,
   140  // return false if this process can't insert this value, because another process has insert the same value.
   141  //
   142  // If the value is in the skip set but not fully linked, this process will wait until it is.
   143  func (s *Int64Set) Add(value int64) bool {
   144  	level := s.randomlevel()
   145  	var preds, succs [maxLevel]*int64Node
   146  	for {
   147  		lFound := s.findNodeAdd(value, &preds, &succs)
   148  		if lFound != -1 { // indicating the value is already in the skip-list
   149  			nodeFound := succs[lFound]
   150  			if !nodeFound.flags.Get(marked) {
   151  				for !nodeFound.flags.Get(fullyLinked) {
   152  					// The node is not yet fully linked, just waits until it is.
   153  				}
   154  				return false
   155  			}
   156  			// If the node is marked, represents some other thread is in the process of deleting this node,
   157  			// we need to add this node in next loop.
   158  			continue
   159  		}
   160  		// Add this node into skip list.
   161  		var (
   162  			highestLocked        = -1 // the highest level being locked by this process
   163  			valid                = true
   164  			pred, succ, prevPred *int64Node
   165  		)
   166  		for layer := 0; valid && layer < level; layer++ {
   167  			pred = preds[layer]   // target node's previous node
   168  			succ = succs[layer]   // target node's next node
   169  			if pred != prevPred { // the node in this layer could be locked by previous loop
   170  				pred.mu.Lock()
   171  				highestLocked = layer
   172  				prevPred = pred
   173  			}
   174  			// valid check if there is another node has inserted into the skip list in this layer during this process.
   175  			// It is valid if:
   176  			// 1. The previous node and next node both are not marked.
   177  			// 2. The previous node's next node is succ in this layer.
   178  			valid = !pred.flags.Get(marked) && (succ == nil || !succ.flags.Get(marked)) && pred.loadNext(layer) == succ
   179  		}
   180  		if !valid {
   181  			unlockInt64(preds, highestLocked)
   182  			continue
   183  		}
   184  
   185  		nn := newInt64Node(value, level)
   186  		for layer := 0; layer < level; layer++ {
   187  			nn.storeNext(layer, succs[layer])
   188  			preds[layer].atomicStoreNext(layer, nn)
   189  		}
   190  		nn.flags.SetTrue(fullyLinked)
   191  		unlockInt64(preds, highestLocked)
   192  		atomic.AddInt64(&s.length, 1)
   193  		return true
   194  	}
   195  }
   196  
   197  func (s *Int64Set) randomlevel() int {
   198  	// Generate random level.
   199  	level := randomLevel()
   200  	// Update highest level if possible.
   201  	for {
   202  		hl := atomic.LoadInt64(&s.highestLevel)
   203  		if int64(level) <= hl {
   204  			break
   205  		}
   206  		if atomic.CompareAndSwapInt64(&s.highestLevel, hl, int64(level)) {
   207  			break
   208  		}
   209  	}
   210  	return level
   211  }
   212  
   213  // Contains check if the value is in the skip set.
   214  func (s *Int64Set) Contains(value int64) bool {
   215  	x := s.header
   216  	for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- {
   217  		nex := x.atomicLoadNext(i)
   218  		for nex != nil && nex.lessthan(value) {
   219  			x = nex
   220  			nex = x.atomicLoadNext(i)
   221  		}
   222  
   223  		// Check if the value already in the skip list.
   224  		if nex != nil && nex.equal(value) {
   225  			return nex.flags.MGet(fullyLinked|marked, fullyLinked)
   226  		}
   227  	}
   228  	return false
   229  }
   230  
   231  // Remove a node from the skip set.
   232  func (s *Int64Set) Remove(value int64) bool {
   233  	var (
   234  		nodeToRemove *int64Node
   235  		isMarked     bool // represents if this operation mark the node
   236  		topLayer     = -1
   237  		preds, succs [maxLevel]*int64Node
   238  	)
   239  	for {
   240  		lFound := s.findNodeRemove(value, &preds, &succs)
   241  		if isMarked || // this process mark this node or we can find this node in the skip list
   242  			lFound != -1 && succs[lFound].flags.MGet(fullyLinked|marked, fullyLinked) && (int(succs[lFound].level)-1) == lFound {
   243  			if !isMarked { // we don't mark this node for now
   244  				nodeToRemove = succs[lFound]
   245  				topLayer = lFound
   246  				nodeToRemove.mu.Lock()
   247  				if nodeToRemove.flags.Get(marked) {
   248  					// The node is marked by another process,
   249  					// the physical deletion will be accomplished by another process.
   250  					nodeToRemove.mu.Unlock()
   251  					return false
   252  				}
   253  				nodeToRemove.flags.SetTrue(marked)
   254  				isMarked = true
   255  			}
   256  			// Accomplish the physical deletion.
   257  			var (
   258  				highestLocked        = -1 // the highest level being locked by this process
   259  				valid                = true
   260  				pred, succ, prevPred *int64Node
   261  			)
   262  			for layer := 0; valid && (layer <= topLayer); layer++ {
   263  				pred, succ = preds[layer], succs[layer]
   264  				if pred != prevPred { // the node in this layer could be locked by previous loop
   265  					pred.mu.Lock()
   266  					highestLocked = layer
   267  					prevPred = pred
   268  				}
   269  				// valid check if there is another node has inserted into the skip list in this layer
   270  				// during this process, or the previous is removed by another process.
   271  				// It is valid if:
   272  				// 1. the previous node exists.
   273  				// 2. no another node has inserted into the skip list in this layer.
   274  				valid = !pred.flags.Get(marked) && pred.loadNext(layer) == succ
   275  			}
   276  			if !valid {
   277  				unlockInt64(preds, highestLocked)
   278  				continue
   279  			}
   280  			for i := topLayer; i >= 0; i-- {
   281  				// Now we own the `nodeToRemove`, no other goroutine will modify it.
   282  				// So we don't need `nodeToRemove.loadNext`
   283  				preds[i].atomicStoreNext(i, nodeToRemove.loadNext(i))
   284  			}
   285  			nodeToRemove.mu.Unlock()
   286  			unlockInt64(preds, highestLocked)
   287  			atomic.AddInt64(&s.length, -1)
   288  			return true
   289  		}
   290  		return false
   291  	}
   292  }
   293  
   294  // Range calls f sequentially for each value present in the skip set.
   295  // If f returns false, range stops the iteration.
   296  func (s *Int64Set) Range(f func(value int64) bool) {
   297  	x := s.header.atomicLoadNext(0)
   298  	for x != nil {
   299  		if !x.flags.MGet(fullyLinked|marked, fullyLinked) {
   300  			x = x.atomicLoadNext(0)
   301  			continue
   302  		}
   303  		if !f(x.value) {
   304  			break
   305  		}
   306  		x = x.atomicLoadNext(0)
   307  	}
   308  }
   309  
   310  // Len return the length of this skip set.
   311  func (s *Int64Set) Len() int {
   312  	return int(atomic.LoadInt64(&s.length))
   313  }