feat: skipmap support LoadOrStoreLazy

collection/skipmap/skipmap.go

@@ -242,6 +242,7 @@ func (s *Int64Map) Load(key int64) (value interface{}, ok bool) {
 
 // LoadAndDelete deletes the value for a key, returning the previous value if any.
 // The loaded result reports whether the key was present.
+// (Modified from Delete)
 func (s *Int64Map) LoadAndDelete(key int64) (value interface{}, loaded bool) {
 	var (
 		nodeToDelete *int64Node
@@ -307,6 +308,7 @@ func (s *Int64Map) LoadAndDelete(key int64) (value interface{}, loaded bool) {
 // LoadOrStore returns the existing value for the key if present.
 // Otherwise, it stores and returns the given value.
 // The loaded result is true if the value was loaded, false if stored.
+// (Modified from Store)
 func (s *Int64Map) LoadOrStore(key int64, value interface{}) (actual interface{}, loaded bool) {
 	level := s.randomlevel()
 	var preds, succs [maxLevel]*int64Node
@@ -360,6 +362,63 @@ func (s *Int64Map) LoadOrStore(key int64, value interface{}) (actual interface{}
 	}
 }
 
+// LoadOrStoreLazy returns the existing value for the key if present.
+// Otherwise, it stores and returns the given value from f, f will only be called once.
+// The loaded result is true if the value was loaded, false if stored.
+// (Modified from LoadOrStore)
+func (s *Int64Map) LoadOrStoreLazy(key int64, f func() interface{}) (actual interface{}, loaded bool) {
+	level := s.randomlevel()
+	var preds, succs [maxLevel]*int64Node
+	for {
+		nodeFound := s.findNode(key, &preds, &succs)
+		if nodeFound != nil { // indicating the key is already in the skip-list
+			if !nodeFound.flags.Get(marked) {
+				// We don't need to care about whether or not the node is fully linked,
+				// just return the value.
+				return nodeFound.loadVal(), true
+			}
+			// If the node is marked, represents some other goroutines is in the process of deleting this node,
+			// we need to add this node in next loop.
+			continue
+		}
+
+		// Add this node into skip list.
+		var (
+			highestLocked        = -1 // the highest level being locked by this process
+			valid                = true
+			pred, succ, prevPred *int64Node
+		)
+		for layer := 0; valid && layer < level; layer++ {
+			pred = preds[layer]   // target node's previous node
+			succ = succs[layer]   // target node's next node
+			if pred != prevPred { // the node in this layer could be locked by previous loop
+				pred.mu.Lock()
+				highestLocked = layer
+				prevPred = pred
+			}
+			// valid check if there is another node has inserted into the skip list in this layer during this process.
+			// It is valid if:
+			// 1. The previous node and next node both are not marked.
+			// 2. The previous node's next node is succ in this layer.
+			valid = !pred.flags.Get(marked) && pred.loadNext(layer) == succ && (succ == nil || !succ.flags.Get(marked))
+		}
+		if !valid {
+			unlockInt64(preds, highestLocked)
+			continue
+		}
+		value := f()
+		nn := newInt64Node(key, value, level)
+		for layer := 0; layer < level; layer++ {
+			nn.storeNext(layer, succs[layer])
+			preds[layer].atomicStoreNext(layer, nn)
+		}
+		nn.flags.SetTrue(fullyLinked)
+		unlockInt64(preds, highestLocked)
+		atomic.AddInt64(&s.length, 1)
+		return value, false
+	}
+}
+
 // Delete deletes the value for a key.
 func (s *Int64Map) Delete(key int64) bool {
 	var (

collection/skipmap/skipmap_test.go

@@ -230,6 +230,35 @@ func TestSkipMap(t *testing.T) {
 	if added != 1 {
 		t.Fatal("Only one LoadAndDelete can successfully get a value")
 	}
+
+	// Correntness 5. (LoadOrStoreLazy)
+	mp = NewInt()
+	tmpmap = NewInt64()
+	samekey = 123
+	added = 0
+	var fcalled int64
+	valuef := func() interface{} {
+		atomic.AddInt64(&fcalled, 1)
+		return fastrand.Int63()
+	}
+	for i := 1; i < 1000; i++ {
+		wg.Add(1)
+		go func() {
+			actual, loaded := mp.LoadOrStoreLazy(samekey, valuef)
+			if !loaded {
+				atomic.AddInt64(&added, 1)
+			}
+			tmpmap.Store(actual.(int64), nil)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	if added != 1 || fcalled != 1 {
+		t.Fatal("only one LoadOrStoreLazy can successfully insert a key and value")
+	}
+	if tmpmap.Len() != 1 {
+		t.Fatal("only one value can be returned from LoadOrStoreLazy")
+	}
 }
 
 func TestSkipMapDesc(t *testing.T) {