github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/structure/maps/skipmap/skipmap.go (about) 1 // Package skipmap is a high-performance, scalable, concurrent-safe map based on skip-list. 2 // In the typical pattern(100000 operations, 90%LOAD 9%STORE 1%DELETE, 8C16T), the skipmap 3 // up to 10x faster than the built-in sync.Map. 4 package skipmap 5 6 import ( 7 "fmt" 8 "strings" 9 "sync" 10 "sync/atomic" 11 "unsafe" 12 13 "github.com/songzhibin97/go-baseutils/base/bcomparator" 14 "github.com/songzhibin97/go-baseutils/structure/maps" 15 ) 16 17 // Assert Map implementation 18 var _ maps.Map[int, any] = (*Map[int, any])(nil) 19 20 // Map represents a map based on skip list in ascending order. 21 type Map[K, V any] struct { 22 header *node[K, V] 23 length int64 24 highestLevel int64 // highest level for now 25 comparator bcomparator.Comparator[K] 26 zeroV V 27 } 28 29 func (s *Map[K, V]) Put(key K, value V) { 30 s.Store(key, value) 31 } 32 33 func (s *Map[K, V]) Get(key K) (value V, found bool) { 34 return s.Load(key) 35 } 36 37 func (s *Map[K, V]) Remove(key K) { 38 s.Delete(key) 39 } 40 41 func (s *Map[K, V]) Keys() []K { 42 ks := make([]K, 0, s.Len()) 43 s.Range(func(k K, v V) bool { 44 ks = append(ks, k) 45 return true 46 }) 47 return ks 48 } 49 50 func (s *Map[K, V]) Empty() bool { 51 return s.Len() == 0 52 } 53 54 func (s *Map[K, V]) Size() int { 55 return s.Len() 56 } 57 58 func (s *Map[K, V]) Clear() { 59 var zerok K 60 var zerov V 61 h := newNode[K, V](zerok, zerov, maxLevel, s.comparator) 62 h.flags.SetTrue(fullyLinked) 63 s.header = h 64 s.highestLevel = defaultHighestLevel 65 } 66 67 func (s *Map[K, V]) Values() []V { 68 vs := make([]V, 0, s.Len()) 69 s.Range(func(k K, v V) bool { 70 vs = append(vs, v) 71 return true 72 }) 73 return vs 74 } 75 76 func (s *Map[K, V]) String() string { 77 bf := strings.Builder{} 78 bf.WriteString("SkipMap\nmap[") 79 s.Range(func(k K, v V) bool { 80 bf.WriteString(fmt.Sprintf("(%v:%v) ", k, v)) 81 return true 82 }) 83 bf.WriteString("]") 84 return bf.String() 85 } 86 87 type node[K, V any] struct { 88 key K 89 value unsafe.Pointer // unsafe.Pointer *V{} 90 next optionalArray // [level]*node 91 mu sync.Mutex 92 flags bitflag 93 level uint32 94 comparator bcomparator.Comparator[K] 95 } 96 97 func newNode[K, V any](key K, value V, level int, comparator bcomparator.Comparator[K]) *node[K, V] { 98 node := &node[K, V]{ 99 key: key, 100 level: uint32(level), 101 comparator: comparator, 102 } 103 node.storeVal(value) 104 if level > op1 { 105 node.next.extra = new([op2]unsafe.Pointer) 106 } 107 return node 108 } 109 110 func (n *node[K, V]) storeVal(value V) { 111 atomic.StorePointer(&n.value, unsafe.Pointer(&value)) 112 } 113 114 func (n *node[K, V]) loadVal() V { 115 return *(*V)(atomic.LoadPointer(&n.value)) 116 } 117 118 func (n *node[K, V]) loadNext(i int) *node[K, V] { 119 return (*node[K, V])(n.next.load(i)) 120 } 121 122 func (n *node[K, V]) storeNext(i int, node *node[K, V]) { 123 n.next.store(i, unsafe.Pointer(node)) 124 } 125 126 func (n *node[K, V]) atomicLoadNext(i int) *node[K, V] { 127 return (*node[K, V])(n.next.atomicLoad(i)) 128 } 129 130 func (n *node[K, V]) atomicStoreNext(i int, node *node[K, V]) { 131 n.next.atomicStore(i, unsafe.Pointer(node)) 132 } 133 134 func (n *node[K, V]) lessthan(key K, comparator bcomparator.Comparator[K]) bool { 135 return comparator(n.key, key) < 0 136 } 137 138 func (n *node[K, V]) equal(key K, comparator bcomparator.Comparator[K]) bool { 139 return comparator(n.key, key) == 0 140 } 141 142 // New return an empty int64 skipmap. 143 func New[K, V any](comparator bcomparator.Comparator[K]) *Map[K, V] { 144 var zerok K 145 var zerov V 146 h := newNode[K, V](zerok, zerov, maxLevel, comparator) 147 h.flags.SetTrue(fullyLinked) 148 return &Map[K, V]{ 149 header: h, 150 highestLevel: defaultHighestLevel, 151 comparator: comparator, 152 } 153 } 154 155 // findNode takes a key and two maximal-height arrays then searches exactly as in a sequential skipmap. 156 // The returned preds and succs always satisfy preds[i] > key >= succs[i]. 157 // (without fullpath, if find the node will return immediately) 158 func (s *Map[K, V]) findNode(key K, preds *[maxLevel]*node[K, V], succs *[maxLevel]*node[K, V]) *node[K, V] { 159 x := s.header 160 for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- { 161 succ := x.atomicLoadNext(i) 162 for succ != nil && succ.lessthan(key, s.comparator) { 163 x = succ 164 succ = x.atomicLoadNext(i) 165 } 166 preds[i] = x 167 succs[i] = succ 168 169 // Check if the key already in the skipmap. 170 if succ != nil && succ.equal(key, s.comparator) { 171 return succ 172 } 173 } 174 return nil 175 } 176 177 // findNodeDelete takes a key and two maximal-height arrays then searches exactly as in a sequential skip-list. 178 // The returned preds and succs always satisfy preds[i] > key >= succs[i]. 179 func (s *Map[K, V]) findNodeDelete(key K, preds *[maxLevel]*node[K, V], succs *[maxLevel]*node[K, V]) int { 180 // lFound represents the index of the first layer at which it found a node. 181 lFound, x := -1, s.header 182 for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- { 183 succ := x.atomicLoadNext(i) 184 for succ != nil && succ.lessthan(key, s.comparator) { 185 x = succ 186 succ = x.atomicLoadNext(i) 187 } 188 preds[i] = x 189 succs[i] = succ 190 191 // Check if the key already in the skip list. 192 if lFound == -1 && succ != nil && succ.equal(key, s.comparator) { 193 lFound = i 194 } 195 } 196 return lFound 197 } 198 199 func unlock[K, V any](preds [maxLevel]*node[K, V], highestLevel int) { 200 var prevPred *node[K, V] 201 for i := highestLevel; i >= 0; i-- { 202 if preds[i] != prevPred { // the node could be unlocked by previous loop 203 preds[i].mu.Unlock() 204 prevPred = preds[i] 205 } 206 } 207 } 208 209 // Store sets the value for a key. 210 func (s *Map[K, V]) Store(key K, value V) { 211 level := s.randomlevel() 212 var preds, succs [maxLevel]*node[K, V] 213 for { 214 nodeFound := s.findNode(key, &preds, &succs) 215 if nodeFound != nil { // indicating the key is already in the skip-list 216 if !nodeFound.flags.Get(marked) { 217 // We don't need to care about whether or not the node is fully linked, 218 // just replace the value. 219 nodeFound.storeVal(value) 220 return 221 } 222 // If the node is marked, represents some other goroutines is in the process of deleting this node, 223 // we need to add this node in next loop. 224 continue 225 } 226 227 // Add this node into skip list. 228 var ( 229 highestLocked = -1 // the highest level being locked by this process 230 valid = true 231 pred, succ, prevPred *node[K, V] 232 ) 233 for layer := 0; valid && layer < level; layer++ { 234 pred = preds[layer] // target node's previous node 235 succ = succs[layer] // target node's next node 236 if pred != prevPred { // the node in this layer could be locked by previous loop 237 pred.mu.Lock() 238 highestLocked = layer 239 prevPred = pred 240 } 241 // valid check if there is another node has inserted into the skip list in this layer during this process. 242 // It is valid if: 243 // 1. The previous node and next node both are not marked. 244 // 2. The previous node's next node is succ in this layer. 245 valid = !pred.flags.Get(marked) && (succ == nil || !succ.flags.Get(marked)) && pred.loadNext(layer) == succ 246 } 247 if !valid { 248 unlock(preds, highestLocked) 249 continue 250 } 251 252 nn := newNode[K, V](key, value, level, s.comparator) 253 for layer := 0; layer < level; layer++ { 254 nn.storeNext(layer, succs[layer]) 255 preds[layer].atomicStoreNext(layer, nn) 256 } 257 nn.flags.SetTrue(fullyLinked) 258 unlock(preds, highestLocked) 259 atomic.AddInt64(&s.length, 1) 260 return 261 } 262 } 263 264 func (s *Map[K, V]) randomlevel() int { 265 // Generate random level. 266 level := randomLevel() 267 // Update highest level if possible. 268 for { 269 hl := atomic.LoadInt64(&s.highestLevel) 270 if int64(level) <= hl { 271 break 272 } 273 if atomic.CompareAndSwapInt64(&s.highestLevel, hl, int64(level)) { 274 break 275 } 276 } 277 return level 278 } 279 280 // Load returns the value stored in the map for a key, or nil if no 281 // value is present. 282 // The ok result indicates whether value was found in the map. 283 func (s *Map[K, V]) Load(key K) (value V, ok bool) { 284 x := s.header 285 for i := int(atomic.LoadInt64(&s.highestLevel)) - 1; i >= 0; i-- { 286 nex := x.atomicLoadNext(i) 287 for nex != nil && nex.lessthan(key, s.comparator) { 288 x = nex 289 nex = x.atomicLoadNext(i) 290 } 291 292 // Check if the key already in the skip list. 293 if nex != nil && nex.equal(key, s.comparator) { 294 if nex.flags.MGet(fullyLinked|marked, fullyLinked) { 295 return nex.loadVal(), true 296 } 297 return s.zeroV, false 298 } 299 } 300 return s.zeroV, false 301 } 302 303 // LoadAndDelete deletes the value for a key, returning the previous value if any. 304 // The loaded result reports whether the key was present. 305 // (Modified from Delete) 306 func (s *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) { 307 var ( 308 nodeToDelete *node[K, V] 309 isMarked bool // represents if this operation mark the node 310 topLayer = -1 311 preds, succs [maxLevel]*node[K, V] 312 ) 313 for { 314 lFound := s.findNodeDelete(key, &preds, &succs) 315 if isMarked || // this process mark this node or we can find this node in the skip list 316 lFound != -1 && succs[lFound].flags.MGet(fullyLinked|marked, fullyLinked) && (int(succs[lFound].level)-1) == lFound { 317 if !isMarked { // we don't mark this node for now 318 nodeToDelete = succs[lFound] 319 topLayer = lFound 320 nodeToDelete.mu.Lock() 321 if nodeToDelete.flags.Get(marked) { 322 // The node is marked by another process, 323 // the physical deletion will be accomplished by another process. 324 nodeToDelete.mu.Unlock() 325 return s.zeroV, false 326 } 327 nodeToDelete.flags.SetTrue(marked) 328 isMarked = true 329 } 330 // Accomplish the physical deletion. 331 var ( 332 highestLocked = -1 // the highest level being locked by this process 333 valid = true 334 pred, succ, prevPred *node[K, V] 335 ) 336 for layer := 0; valid && (layer <= topLayer); layer++ { 337 pred, succ = preds[layer], succs[layer] 338 if pred != prevPred { // the node in this layer could be locked by previous loop 339 pred.mu.Lock() 340 highestLocked = layer 341 prevPred = pred 342 } 343 // valid check if there is another node has inserted into the skip list in this layer 344 // during this process, or the previous is deleted by another process. 345 // It is valid if: 346 // 1. the previous node exists. 347 // 2. no another node has inserted into the skip list in this layer. 348 valid = !pred.flags.Get(marked) && pred.loadNext(layer) == succ 349 } 350 if !valid { 351 unlock(preds, highestLocked) 352 continue 353 } 354 for i := topLayer; i >= 0; i-- { 355 // Now we own the `nodeToDelete`, no other goroutine will modify it. 356 // So we don't need `nodeToDelete.loadNext` 357 preds[i].atomicStoreNext(i, nodeToDelete.loadNext(i)) 358 } 359 nodeToDelete.mu.Unlock() 360 unlock(preds, highestLocked) 361 atomic.AddInt64(&s.length, -1) 362 return nodeToDelete.loadVal(), true 363 } 364 return s.zeroV, false 365 } 366 } 367 368 // LoadOrStore returns the existing value for the key if present. 369 // Otherwise, it stores and returns the given value. 370 // The loaded result is true if the value was loaded, false if stored. 371 // (Modified from Store) 372 func (s *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) { 373 var ( 374 level int 375 preds, succs [maxLevel]*node[K, V] 376 hl = int(atomic.LoadInt64(&s.highestLevel)) 377 ) 378 for { 379 nodeFound := s.findNode(key, &preds, &succs) 380 if nodeFound != nil { // indicating the key is already in the skip-list 381 if !nodeFound.flags.Get(marked) { 382 // We don't need to care about whether or not the node is fully linked, 383 // just return the value. 384 return nodeFound.loadVal(), true 385 } 386 // If the node is marked, represents some other goroutines is in the process of deleting this node, 387 // we need to add this node in next loop. 388 continue 389 } 390 391 // Add this node into skip list. 392 var ( 393 highestLocked = -1 // the highest level being locked by this process 394 valid = true 395 pred, succ, prevPred *node[K, V] 396 ) 397 if level == 0 { 398 level = s.randomlevel() 399 if level > hl { 400 // If the highest level is updated, usually means that many goroutines 401 // are inserting items. Hopefully we can find a better path in next loop. 402 // TODO(zyh): consider filling the preds if s.header[level].next == nil, 403 // but this strategy's performance is almost the same as the existing method. 404 continue 405 } 406 } 407 for layer := 0; valid && layer < level; layer++ { 408 pred = preds[layer] // target node's previous node 409 succ = succs[layer] // target node's next node 410 if pred != prevPred { // the node in this layer could be locked by previous loop 411 pred.mu.Lock() 412 highestLocked = layer 413 prevPred = pred 414 } 415 // valid check if there is another node has inserted into the skip list in this layer during this process. 416 // It is valid if: 417 // 1. The previous node and next node both are not marked. 418 // 2. The previous node's next node is succ in this layer. 419 valid = !pred.flags.Get(marked) && (succ == nil || !succ.flags.Get(marked)) && pred.loadNext(layer) == succ 420 } 421 if !valid { 422 unlock(preds, highestLocked) 423 continue 424 } 425 426 nn := newNode(key, value, level, s.comparator) 427 for layer := 0; layer < level; layer++ { 428 nn.storeNext(layer, succs[layer]) 429 preds[layer].atomicStoreNext(layer, nn) 430 } 431 nn.flags.SetTrue(fullyLinked) 432 unlock(preds, highestLocked) 433 atomic.AddInt64(&s.length, 1) 434 return value, false 435 } 436 } 437 438 // LoadOrStoreLazy returns the existing value for the key if present. 439 // Otherwise, it stores and returns the given value from f, f will only be called once. 440 // The loaded result is true if the value was loaded, false if stored. 441 // (Modified from LoadOrStore) 442 func (s *Map[K, V]) LoadOrStoreLazy(key K, f func() V) (actual V, loaded bool) { 443 var ( 444 level int 445 preds, succs [maxLevel]*node[K, V] 446 hl = int(atomic.LoadInt64(&s.highestLevel)) 447 ) 448 for { 449 nodeFound := s.findNode(key, &preds, &succs) 450 if nodeFound != nil { // indicating the key is already in the skip-list 451 if !nodeFound.flags.Get(marked) { 452 // We don't need to care about whether or not the node is fully linked, 453 // just return the value. 454 return nodeFound.loadVal(), true 455 } 456 // If the node is marked, represents some other goroutines is in the process of deleting this node, 457 // we need to add this node in next loop. 458 continue 459 } 460 461 // Add this node into skip list. 462 var ( 463 highestLocked = -1 // the highest level being locked by this process 464 valid = true 465 pred, succ, prevPred *node[K, V] 466 ) 467 if level == 0 { 468 level = s.randomlevel() 469 if level > hl { 470 // If the highest level is updated, usually means that many goroutines 471 // are inserting items. Hopefully we can find a better path in next loop. 472 // TODO(zyh): consider filling the preds if s.header[level].next == nil, 473 // but this strategy's performance is almost the same as the existing method. 474 continue 475 } 476 } 477 for layer := 0; valid && layer < level; layer++ { 478 pred = preds[layer] // target node's previous node 479 succ = succs[layer] // target node's next node 480 if pred != prevPred { // the node in this layer could be locked by previous loop 481 pred.mu.Lock() 482 highestLocked = layer 483 prevPred = pred 484 } 485 // valid check if there is another node has inserted into the skip list in this layer during this process. 486 // It is valid if: 487 // 1. The previous node and next node both are not marked. 488 // 2. The previous node's next node is succ in this layer. 489 valid = !pred.flags.Get(marked) && pred.loadNext(layer) == succ && (succ == nil || !succ.flags.Get(marked)) 490 } 491 if !valid { 492 unlock(preds, highestLocked) 493 continue 494 } 495 value := f() 496 nn := newNode(key, value, level, s.comparator) 497 for layer := 0; layer < level; layer++ { 498 nn.storeNext(layer, succs[layer]) 499 preds[layer].atomicStoreNext(layer, nn) 500 } 501 nn.flags.SetTrue(fullyLinked) 502 unlock(preds, highestLocked) 503 atomic.AddInt64(&s.length, 1) 504 return value, false 505 } 506 } 507 508 // Delete deletes the value for a key. 509 func (s *Map[K, V]) Delete(key K) bool { 510 var ( 511 nodeToDelete *node[K, V] 512 isMarked bool // represents if this operation mark the node 513 topLayer = -1 514 preds, succs [maxLevel]*node[K, V] 515 ) 516 for { 517 lFound := s.findNodeDelete(key, &preds, &succs) 518 if isMarked || // this process mark this node or we can find this node in the skip list 519 lFound != -1 && succs[lFound].flags.MGet(fullyLinked|marked, fullyLinked) && (int(succs[lFound].level)-1) == lFound { 520 if !isMarked { // we don't mark this node for now 521 nodeToDelete = succs[lFound] 522 topLayer = lFound 523 nodeToDelete.mu.Lock() 524 if nodeToDelete.flags.Get(marked) { 525 // The node is marked by another process, 526 // the physical deletion will be accomplished by another process. 527 nodeToDelete.mu.Unlock() 528 return false 529 } 530 nodeToDelete.flags.SetTrue(marked) 531 isMarked = true 532 } 533 // Accomplish the physical deletion. 534 var ( 535 highestLocked = -1 // the highest level being locked by this process 536 valid = true 537 pred, succ, prevPred *node[K, V] 538 ) 539 for layer := 0; valid && (layer <= topLayer); layer++ { 540 pred, succ = preds[layer], succs[layer] 541 if pred != prevPred { // the node in this layer could be locked by previous loop 542 pred.mu.Lock() 543 highestLocked = layer 544 prevPred = pred 545 } 546 // valid check if there is another node has inserted into the skip list in this layer 547 // during this process, or the previous is deleted by another process. 548 // It is valid if: 549 // 1. the previous node exists. 550 // 2. no another node has inserted into the skip list in this layer. 551 valid = !pred.flags.Get(marked) && pred.atomicLoadNext(layer) == succ 552 } 553 if !valid { 554 unlock(preds, highestLocked) 555 continue 556 } 557 for i := topLayer; i >= 0; i-- { 558 // Now we own the `nodeToDelete`, no other goroutine will modify it. 559 // So we don't need `nodeToDelete.loadNext` 560 preds[i].atomicStoreNext(i, nodeToDelete.loadNext(i)) 561 } 562 nodeToDelete.mu.Unlock() 563 unlock(preds, highestLocked) 564 atomic.AddInt64(&s.length, -1) 565 return true 566 } 567 return false 568 } 569 } 570 571 // Range calls f sequentially for each key and value present in the skipmap. 572 // If f returns false, range stops the iteration. 573 // 574 // Range does not necessarily correspond to any consistent snapshot of the Map's 575 // contents: no key will be visited more than once, but if the value for any key 576 // is stored or deleted concurrently, Range may reflect any mapping for that key 577 // from any point during the Range call. 578 func (s *Map[K, V]) Range(f func(key K, value V) bool) { 579 x := s.header.atomicLoadNext(0) 580 for x != nil { 581 if !x.flags.MGet(fullyLinked|marked, fullyLinked) { 582 x = x.atomicLoadNext(0) 583 continue 584 } 585 if !f(x.key, x.loadVal()) { 586 break 587 } 588 x = x.atomicLoadNext(0) 589 } 590 } 591 592 // Len return the length of this skipmap. 593 func (s *Map[K, V]) Len() int { 594 return int(atomic.LoadInt64(&s.length)) 595 }