lockfree-treiber-stack

Answer

A lock-free Treiber stack implementing hazard pointers for safe memory reclamation. The core algorithm:

Push – allocate a new Node, point it at the current top, then CAS top to the new node. No hazard pointer needed since the new node is private until published.

Pop – the hazard-pointer protocol has four steps: 1. Load top and immediately announce it in the goroutine's hazard-pointer slot 2. Validate that top hasn't changed (closes the race between the load and the announcement) 3. Read next (safe because the hazard pointer prevents reclamation of top) 4. CAS top from the observed node to its successor; on success, clear the hazard pointer and retire the node

Reclamation – each goroutine keeps a small per-thread retired list. When it reaches retireLimit (3), the goroutine scans all 128 hazard pointer slots globally; nodes still referenced by any slot stay in the retired list, the rest are dropped and become eligible for Go's GC.

ABA prevention – because Push always allocates fresh nodes and never reuses retired ones, a node protected by a hazard pointer (held by another goroutine that loaded it as top) cannot be recycled onto the stack. The retired list maintains a live reference, preventing both premature reclamation and address reuse.

package main

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

// Node in the lock-free Treiber stack.
type Node struct {
    value interface{}
    next  atomic.Pointer[Node]
}

// threadState holds per-goroutine hazard-pointer slot and retired list.
type threadState struct {
    slot    int
    retired []*Node
}

const (
    hpSlots     = 128   // enough for 128 concurrent goroutines
    retireLimit = 3     // trigger reclamation scan every N retirements
)

// Stack implements a lock-free Treiber stack with hazard pointers.
type Stack struct {
    top     atomic.Pointer[Node]
    slotsMu sync.Mutex
    slots   [hpSlots]atomic.Pointer[Node]
    tls     sync.Map       // goid → *threadState
    idGen   atomic.Int32
}

func goid() int64 {
    var buf [64]byte
    n := runtime.Stack(buf[:], false)
    var id int64
    for i := 10; i < n; i++ {
        if buf[i] == ' ' { break }
        id = id*10 + int64(buf[i]-'0')
    }
    return id
}

func NewStack() *Stack { return &Stack{} }

func (s *Stack) getState() *threadState {
    gid := goid()
    if v, ok := s.tls.Load(gid); ok {
        return v.(*threadState)
    }
    s.slotsMu.Lock()
    defer s.slotsMu.Unlock()
    if v, ok := s.tls.Load(gid); ok {
        return v.(*threadState)
    }
    idx := int(s.idGen.Add(1)-1) % hpSlots
    ts := &threadState{slot: idx, retired: make([]*Node, 0, retireLimit)}
    s.tls.Store(gid, ts)
    return ts
}

// Push – lock-free, no hazard pointer needed (node is private until CAS).
func (s *Stack) Push(value interface{}) {
    node := &Node{value: value}
    for {
        top := s.top.Load()
        node.next.Store(top)
        if s.top.CompareAndSwap(top, node) {
            return
        }
        runtime.Gosched()
    }
}

// Pop – lock-free with hazard-pointer protection.
func (s *Stack) Pop() (interface{}, bool) {
    ts := s.getState()
    for {
        top := s.top.Load()
        if top == nil {
            s.slots[ts.slot].Store(nil)
            return nil, false
        }
        // 1. Announce in hazard pointer
        s.slots[ts.slot].Store(top)
        // 2. Validate
        if s.top.Load() != top {
            continue
        }
        // 3. Read successor
        next := top.next.Load()
        // 4. CAS
        if s.top.CompareAndSwap(top, next) {
            s.slots[ts.slot].Store(nil)
            s.retire(ts, top)
            return top.value, true
        }
        runtime.Gosched()
    }
}

func (s *Stack) retire(ts *threadState, node *Node) {
    ts.retired = append(ts.retired, node)
    if len(ts.retired) >= retireLimit {
        s.scan(ts)
    }
}

func (s *Stack) scan(ts *threadState) {
    if len(ts.retired) == 0 { return }
    protected := make(map[*Node]bool)
    for i := 0; i < hpSlots; i++ {
        if p := s.slots[i].Load(); p != nil {
            protected[p] = true
        }
    }
    kept := ts.retired[:0]
    for _, node := range ts.retired {
        if protected[node] {
            kept = append(kept, node)
        }
    }
    ts.retired = kept
}