v8.run/go/exp@v0.0.26-0.20230226010534-afcdbd3f782d/runtime2/mlocal.go (about)

     1  package runtime2
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"sync/atomic"
     7  )
     8  
     9  type MLocal[T any] struct {
    10  	states   atomic.Pointer[[]T]
    11  	resizing atomic.Bool
    12  	mu       sync.Mutex
    13  }
    14  
    15  func NewMLocal[T any]() *MLocal[T] {
    16  	m := new(MLocal[T])
    17  	arr := make([]T, runtime.GOMAXPROCS(-1)*4)
    18  	m.states.Store(&arr)
    19  	return m
    20  }
    21  
    22  func (ml *MLocal[T]) resize(ns int) {
    23  	ml.mu.Lock()
    24  	ml.resizing.Store(true)
    25  	s0 := ml.states.Load()
    26  	if len(*s0) < ns {
    27  		ml.resizing.Store(false)
    28  		ml.mu.Unlock()
    29  		return
    30  	}
    31  	s1 := make([]T, ns)
    32  	copy(s1, *s0)
    33  	ml.states.Store(&s1)
    34  	ml.resizing.Store(false)
    35  	ml.mu.Unlock()
    36  }
    37  
    38  func (ml *MLocal[T]) Get() T {
    39  	for ml.resizing.Load() {
    40  		runtime.Gosched()
    41  	}
    42  	states := ml.states.Load()
    43  
    44  	mid := ProcPin()
    45  	if mid >= len(*states) {
    46  		ml.resize(mid * 2)
    47  	}
    48  	val := (*states)[mid]
    49  	ProcUnpin()
    50  
    51  	return val
    52  }
    53  
    54  func (ml *MLocal[T]) Set(v T) {
    55  	for ml.resizing.Load() {
    56  		runtime.Gosched()
    57  	}
    58  	states := ml.states.Load()
    59  
    60  	mid := ProcPin()
    61  	if mid >= len(*states) {
    62  		ml.resize(mid * 2)
    63  	}
    64  	(*states)[mid] = v
    65  	ProcUnpin()
    66  }