什麼是原子操作？深入淺析go中的原子操作

func TestAtomic(t *testing.T) {
	var sum = 0
	var wg sync.WaitGroup
	wg.Add(1000)

	// 啟動 1000 個協程，每個協程對 sum 做加法操作
	for i := 0; i < 1000; i++ {
		go func() {
			defer wg.Done()
			sum++
		}()
	}

	// 等待所有的協程都執行完畢
	wg.Wait()
	fmt.Println(sum) // 這裡輸出多少呢？
}

func TestAtomic1(t *testing.T) {
	// 將 sum 的型別改成 int32，因為原子操作只能針對 int32、int64、uint32、uint64、uintptr 這幾種型別
	var sum int32 = 0
	var wg sync.WaitGroup
	wg.Add(1000)

    // 啟動 1000 個協程，每個協程對 sum 做加法操作
	for i := 0; i < 1000; i++ {
		go func() {
			defer wg.Done()
			// 將 sum++ 改成下面這樣
			atomic.AddInt32(&sum, 1)
		}()
	}

	wg.Wait()
	fmt.Println(sum) // 輸出 1000
}

func AddInt32(addr *int32, delta int32) (new int32)
func AddInt64(addr *int64, delta int64) (new int64)
func AddUint32(addr *uint32, delta uint32) (new uint32)
func AddUint64(addr *uint64, delta uint64) (new uint64)
func AddUintptr(addr *uintptr, delta uintptr) (new uintptr)

func CompareAndSwapInt32(addr *int32, old, new int32) (swapped bool)
func CompareAndSwapInt64(addr *int64, old, new int64) (swapped bool)
func CompareAndSwapUint32(addr *uint32, old, new uint32) (swapped bool)
func CompareAndSwapUint64(addr *uint64, old, new uint64) (swapped bool)
func CompareAndSwapUintptr(addr *uintptr, old, new uintptr) (swapped bool)
func CompareAndSwapPointer(addr *unsafe.Pointer, old, new unsafe.Pointer) (swapped bool)

func LoadInt32(addr *int32) (val int32)
func LoadInt64(addr *int64) (val int64)
func LoadUint32(addr *uint32) (val uint32)
func LoadUint64(addr *uint64) (val uint64)
func LoadUintptr(addr *uintptr) (val uintptr)
func LoadPointer(addr *unsafe.Pointer) (val unsafe.Pointer)

func StoreInt32(addr *int32, val int32)
func StoreInt64(addr *int64, val int64)
func StoreUint32(addr *uint32, val uint32)
func StoreUint64(addr *uint64, val uint64)
func StoreUintptr(addr *uintpre, val uintptr)
func StorePointer(addr *unsafe.Pointer, val unsafe.Pointer)

func SwapInt32(addr *int32, new int32) (old int32)
func SwapInt64(addr *int64, new int64) (old int64)
func SwapUint32(addr *uint32, new uint32) (old uint32)
func SwapUint64(addr *uint64, new uint64) (old uint64)
func SwapUintptr(addr *uintptr, new uintptr) (old uintptr)
func SwapPointer(addr *unsafe.Pointer, new unsafe.Pointer) (old unsafe.Pointer)

// Value 提供一致型別值的原子載入和儲存。
type Value struct {
	v any
}

// Load 返回由最近的 Store 設定的值。
func (v *Value) Load() (val any) {
	// atomic.Value 轉換為 efaceWords
	vp := (*efaceWords)(unsafe.Pointer(v))

	// 判斷 atomic.Value 的型別
	typ := LoadPointer(&vp.typ)
	// 第一次 Store 還沒有完成，直接返回 nil
	if typ == nil || typ == unsafe.Pointer(&firstStoreInProgress) {
		// firstStoreInProgress 是一個特殊的變數，儲存到 typ 中用來表示第一次 Store 還沒有完成
		return nil
	}

	// 獲取 atomic.Value 的值
	data := LoadPointer(&vp.data)
	// 將 val 轉換為 efaceWords 型別
	vlp := (*efaceWords)(unsafe.Pointer(&val))
	// 分別賦值給 val 的 typ 和 data
	vlp.typ = typ
	vlp.data = data
	return
}

// 表示一個 interface{}/any 型別
type efaceWords struct {
	typ  unsafe.Pointer
	data unsafe.Pointer
}

// Store 將 Value 的值設定為 val。
func (v *Value) Store(val any) {
	// 不能儲存 nil 值
	if val == nil {
		panic("sync/atomic: store of nil value into Value")
	}
	// atomic.Value 轉換為 efaceWords
	vp := (*efaceWords)(unsafe.Pointer(v))
	// val 轉換為 efaceWords
	vlp := (*efaceWords)(unsafe.Pointer(&val))
	
	// 自旋進行原子操作，這個過程不會很久，開銷相比互斥鎖小
	for {
		// LoadPointer 可以保證獲取到的是最新的
		typ := LoadPointer(&vp.typ)
		// 第一次 store 的時候 typ 還是 nil，說明是第一次 store
		if typ == nil {
			// 嘗試開始第一次 Store。
			// 禁用搶佔，以便其他 goroutines 可以自旋等待完成。
			// （如果允許搶佔，那麼其他 goroutine 自旋等待的時間可能會比較長，因為可能會需要進行協程排程。）
			runtime_procPin()
			// 搶佔失敗，意味著有其他 goroutine 成功 store 了，允許搶佔，再次嘗試 Store
			// 這也是一個原子操作。
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// 完成第一次 store
			// 因為有 firstStoreInProgress 標識的保護，所以下面的兩個原子操作是安全的。
			StorePointer(&vp.data, vlp.data) // 儲存值（原子操作）
			StorePointer(&vp.typ, vlp.typ)   // 儲存型別（原子操作）
			runtime_procUnpin()              // 允許搶佔
			return
		}

		// 另外一個 goroutine 正在進行第一次 Store。自旋等待。
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			continue
		}

		// 第一次 Store 已經完成了，下面不是第一次 Store 了。
		// 需要檢查當前 Store 的型別跟第一次 Store 的型別是否一致，不一致就 panic。
		if typ != vlp.typ {
			panic("sync/atomic: store of inconsistently typed value into Value")
		}

		// 後續的 Store 只需要 Store 值部分就可以了。
		// 因為 atomic.Value 只能儲存一種型別的值。
		StorePointer(&vp.data, vlp.data)
		return
	}
}

// Swap 將 Value 的值設定為 new 並返回舊值。
func (v *Value) Swap(new any) (old any) {
	// 不能儲存 nil 值
	if new == nil {
		panic("sync/atomic: swap of nil value into Value")
	}

	// atomic.Value 轉換為 efaceWords
	vp := (*efaceWords)(unsafe.Pointer(v))
	// new 轉換為 efaceWords
	np := (*efaceWords)(unsafe.Pointer(&new))
	
	// 自旋進行原子操作，這個過程不會很久，開銷相比互斥鎖小
	for {
		// 下面這部分程式碼跟 Store 一樣，不細說了。
		// 這部分程式碼是進行第一次儲存的程式碼。
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return nil
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			continue
		}
		if typ != np.typ {
			panic("sync/atomic: swap of inconsistently typed value into Value")
		}

		// ---- 下面是 Swap 的特有邏輯 ----
		// op 是返回值
		op := (*efaceWords)(unsafe.Pointer(&old))
		// 返回舊的值
		op.typ, op.data = np.typ, SwapPointer(&vp.data, np.data)
		return old
	}
}

// CompareAndSwap 比較並交換。
func (v *Value) CompareAndSwap(old, new any) (swapped bool) {
	// 注意：old 是可以為 nil 的，new 不能為 nil。
	// old 是 nil 表示是第一次進行 Store 操作。
	if new == nil {
		panic("sync/atomic: compare and swap of nil value into Value")
	}

	// atomic.Value 轉換為 efaceWords
	vp := (*efaceWords)(unsafe.Pointer(v))
	// new 轉換為 efaceWords
	np := (*efaceWords)(unsafe.Pointer(&new))
	// old 轉換為 efaceWords
	op := (*efaceWords)(unsafe.Pointer(&old))

	// old 和 new 型別必須一致，且不能為 nil
	if op.typ != nil && np.typ != op.typ {
		panic("sync/atomic: compare and swap of inconsistently typed values")
	}

	// 自旋進行原子操作，這個過程不會很久，開銷相比互斥鎖小
	for {
		// LoadPointer 可以保證獲取到的 typ 是最新的
		typ := LoadPointer(&vp.typ)
		if typ == nil { // atomic.Value 是 nil，還沒 Store 過
			// 準備進行第一次 Store，但是傳遞進來的 old 不是 nil，compare 這一步就失敗了。直接返回 false
			if old != nil {
				return false
			}

			// 下面這部分程式碼跟 Store 一樣，不細說了。 
			// 這部分程式碼是進行第一次儲存的程式碼。
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return true
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			continue
		}
		if typ != np.typ {
			panic("sync/atomic: compare and swap of inconsistently typed value into Value")
		}

		// 通過執行時相等性檢查比較舊版本和當前版本。
		// 這允許對值型別進行比較，這是包函數所沒有的。
		// 下面的 CompareAndSwapPointer 僅確保 vp.data 自 LoadPointer 以來沒有更改。
		data := LoadPointer(&vp.data)
		var i any
		(*efaceWords)(unsafe.Pointer(&i)).typ = typ
		(*efaceWords)(unsafe.Pointer(&i)).data = data
		if i != old { // atomic.Value 跟 old 不相等
			return false
		}
		// 只做 val 部分的 cas 操作
		return CompareAndSwapPointer(&vp.data, data, np.data)
	}
}

// An Int32 is an atomic int32. The zero value is zero.
type Int32 struct {
	_ noCopy
	v int32
}

// Load atomically loads and returns the value stored in x.
func (x *Int32) Load() int32 { return LoadInt32(&x.v) }

// Store atomically stores val into x.
func (x *Int32) Store(val int32) { StoreInt32(&x.v, val) }

// Swap atomically stores new into x and returns the previous value.
func (x *Int32) Swap(new int32) (old int32) { return SwapInt32(&x.v, new) }

// CompareAndSwap executes the compare-and-swap operation for x.
func (x *Int32) CompareAndSwap(old, new int32) (swapped bool) {
	return CompareAndSwapInt32(&x.v, old, new)
}

// 系統資訊 cpu: Intel(R) Core(TM) i7-8700 CPU @ 3.20GHz
// 10.13 ns/op
func BenchmarkMutex(b *testing.B) {
   var mu sync.Mutex

   for i := 0; i < b.N; i++ {
      mu.Lock()
      mu.Unlock()
   }
}

// 5.849 ns/op
func BenchmarkAtomic(b *testing.B) {
   var sum atomic.Uint64

   for i := 0; i < b.N; i++ {
      sum.Add(uint64(1))
   }
}

// 刪除 entry
func (e *entry) delete() (value any, ok bool) {
	for {
		p := e.p.Load()
		// 已經被刪除了，不需要再刪除
		if p == nil || p == expunged {
			return nil, false
		}
		// 刪除成功
		if e.p.CompareAndSwap(p, nil) {
			return *p, true
		}
	}
}

// 如果條目尚未刪除，trySwap 將交換一個值。
func (e *entry) trySwap(i *any) (*any, bool) {
	for {
		p := e.p.Load()
		// 已經被刪除了
		if p == expunged {
			return nil, false
		}
		// swap 成功
		if e.p.CompareAndSwap(p, i) {
			return p, true
		}
	}
}

// WaitGroup 的 Add 函數：高 32 位加上 delta
state := wg.state.Add(uint64(delta) << 32)

// WaitGroup 的 Wait 函數：低 32 位加 1
// 等待者的數量加 1
wg.state.CompareAndSwap(state, state+1)

func TestCas(t *testing.T) {
	var sum int32 = 0
	var wg sync.WaitGroup
	wg.Add(1000)

	for i := 0; i < 1000; i++ {
		go func() {
			defer wg.Done()
			// 這一行是有可能會失敗的
			atomic.CompareAndSwapInt32(&sum, sum, sum+1)
		}()
	}

	wg.Wait()
	fmt.Println(sum) // 不是 1000
}

func TestCas(t *testing.T) {
	var sum int32 = 0
	var wg sync.WaitGroup
	wg.Add(1000)

	for i := 0; i < 1000; i++ {
		go func() {
			defer wg.Done()
			
			// cas 失敗的時候，重新獲取 sum 的值進行計算。
			// cas 成功則返回。
			for {
				if atomic.CompareAndSwapInt32(&sum, sum, sum+1) {
					return
				}
			}
		}()
	}

	wg.Wait()
	fmt.Println(sum)
}