Lock-free, fast GOMAXPROCS(0)

gomaxprocs.go

package xruntime

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

var gmp atomic.Int32

func GOMAXPROCS() int {
	if n := gmp.Load(); n != 0 {
		return int(n)
	}
	return gomaxprocsSlow()
}

func gomaxprocsSlow() int {
	n := runtime.GOMAXPROCS(0)
	if gmp.CompareAndSwap(0, int32(n)) {
		every(1*time.Second, func() {
			n := runtime.GOMAXPROCS(0)
			if n > 0 && int(gmp.Load()) != n {
				gmp.Store(int32(n))
			}
		})
	}
	return n
}


func every(d time.Duration, f func()) func() {
	s := struct {
		mu   sync.Mutex
		stop bool
		t    *time.Timer

		lastPoll time.Time
		fn       func()
		f        func()
		d        time.Duration
	}{
		f: f,
		d: d,
		lastPoll: time.Now(),
	}

	s.fn = func() {
		s.mu.Lock()
		if s.stop {
			s.mu.Unlock()
			return
		}
		s.lastPoll = s.lastPoll.Add(s.d)
		s.t = time.AfterFunc(time.Until(s.lastPoll), s.fn)
		s.mu.Unlock()

		s.f()
	}
	s.lastPoll = s.lastPoll.Add(s.d)
	s.t = time.AfterFunc(time.Until(s.lastPoll), s.fn)

	return func() {
		s.mu.Lock()
		defer s.mu.Unlock()
		if s.stop {
			return
		}
		s.stop = true
		s.t.Stop()
		s.t = nil
		s.f = nil
		s.fn = nil
	}
}