how to implement atomic add but not exceed x?

user_3997070 0
Viewed 94

i want to implement add below in pure atomic operation

var num int
func add(max int) int {
  if num < max {
    num++
  }
  return num
}

i have try one edition

func add(max int64) int64 {
    for {
        old := atomic.LoadInt64(&x)
        if old + 1 < max {
            if atomic.CompareAndSwapInt64(&x, old, old+1) {
                return old+1
            }
        } else {
            return old
        }
    }
}

however, i guess there might be a better solution with less fail chance and avoid dead loop

go atomic
2 Answers

Here's a revised version of your algorithm with a multiple goroutine test. It passes the test with and without the Go race detector.

add.go:

package main

import (
    "fmt"
    "runtime"
    "sync"
    "sync/atomic"
)

var x int64 = -42

func add(max int64) int64 {
    for {
        old := atomic.LoadInt64(&x)
        if old >= max {
            return old
        }
        new := old + 1
        if atomic.CompareAndSwapInt64(&x, old, new) {
            return new
        }
    }
}

func main() {
    const max = 123456
    fmt.Println("max:", max)
    fmt.Println("x:  ", x)
    var wg sync.WaitGroup
    procs := runtime.GOMAXPROCS(0)
    for i := 0; i < procs; i++ {
        wg.Add(1)
        go func(max int64) {
            defer wg.Done()
            for {
                if add(max) >= max {
                    return
                }
            }
        }(max)
    }
    wg.Wait()
    fmt.Println("x:  ", x)

}

https://go.dev/play/p/r-qsnyI7tqv

$ go build -race add.go && ./add
max: 123456
x:   -42
x:   123456
$ go build add.go && ./add
max: 123456
x:   -42
x:   123456
user_164658020

The easiest (but not the most efficient) solution would be to use a mutex:

var (
    mu  sync.Mutex
    num int64
)

func add(max int64) int64 {
    mu.Lock()
    defer mu.Unlock()

    if num < max {
        num++
    }
    return num
}
user_17055980
Related