go 无锁(乐观锁)队列的实现

小灰灰 2022-11-16 13:53 259阅读 0赞

采用轻量级的CAS(Compare and swap)乐观锁，实现安全的队列，代码如下：

/** * 线程安全的队列，使用轻量级的 CAS 锁 */
    package cas_queue
    
    import (
    	"errors"
    	"fmt"
    	"runtime"
    	"sync/atomic"
    )
    
    type casCache struct { 
    	putNo uint32
    	getNo uint32
    	value interface{ }
    }
    
    // lock free queue
    type CASQueue struct { 
    	capacity uint32
    	capMod    uint32
    	putPos    uint32
    	getPos    uint32
    	cache     []casCache
    }
    
    func NewCASQueue(capacity uint32) *CASQueue { 
    	q := new(CASQueue)
    	q.capacity = minQuantity(capacity)
    	q.capMod = q.capacity - 1
    	q.putPos = 0
    	q.getPos = 0
    	q.cache = make([]casCache, q.capacity)
    	for i := range q.cache { 
    		cache := &q.cache[i]
    		cache.getNo = uint32(i)
    		cache.putNo = uint32(i)
    	}
    	cache := &q.cache[0]
    	cache.getNo = q.capacity
    	cache.putNo = q.capacity
    	return q
    }
    
    func (q *CASQueue) String() string { 
    	getPos := atomic.LoadUint32(&q.getPos)
    	putPos := atomic.LoadUint32(&q.putPos)
    	return fmt.Sprintf("Queue{capacity: %v, capMod: %v, putPos: %v, getPos: %v}",
    		q.capacity, q.capMod, putPos, getPos)
    }
    
    func (q *CASQueue) getCapacity() uint32 { 
    	return q.capacity
    }
    
    /** * 获取当前元素个数 */
    func (q *CASQueue) Quantity() uint32 { 
    	var putPos, getPos uint32
    	var quantity uint32
    	getPos = atomic.LoadUint32(&q.getPos)
    	putPos = atomic.LoadUint32(&q.putPos)
    
    	if putPos >= getPos { 
    		quantity = putPos - getPos
    	} else { 
    		quantity = q.capMod + (putPos - getPos)
    	}
    
    	return quantity
    }
    
    /** * put queue functions * ok: 如果成功添加，则为true，反之同一时间有多个线程put导致写入失败或者队列长度不够，则返回false * quantity: 返回代表队列的元素个数，如果大于等于capMod - 1则表示空间满了 */
    func (q *CASQueue) putMayFail(val interface{ }) (ok bool, quantity uint32) { 
    	var putPos, putPosNew, getPos, posCnt uint32
    	var cache *casCache
    	capMod := q.capMod
    
    	getPos = atomic.LoadUint32(&q.getPos)
    	putPos = atomic.LoadUint32(&q.putPos)
    
    	if putPos >= getPos { 
    		posCnt = putPos - getPos
    	} else { 
    		posCnt = capMod + (putPos - getPos)
    	}
    
    	// 空间不足
    	if posCnt >= capMod-1 { 
    		runtime.Gosched()
    		return false, posCnt
    	}
    
    	putPosNew = putPos + 1
    	if !atomic.CompareAndSwapUint32(&q.putPos, putPos, putPosNew) { 
    		runtime.Gosched()
    		return false, posCnt
    	}
    
    	cache = &q.cache[putPosNew & capMod]
    
    	for { 
    		getNo := atomic.LoadUint32(&cache.getNo)
    		putNo := atomic.LoadUint32(&cache.putNo)
    		if putPosNew == putNo && getNo == putNo { 
    			cache.value = val
    			atomic.AddUint32(&cache.putNo, q.capacity)
    			return true, posCnt + 1
    		} else { 
    			runtime.Gosched()
    		}
    	}
    }
    
    /** * 添加一个元素到队列，如果队列满了则报错 */
    func (q *CASQueue) Put(val interface{ }) error { 
    	var ok bool
    	var quantity uint32
    	for !ok {  // 写入失败，没拿到CAS锁，则继续写入
    		ok, quantity = q.putMayFail(val)
    		// 队列长度不够了，则直接返回错误
    		if quantity >= q.capMod - 1 { 
    			errMsg := fmt.Sprintf("queue almost overflow, the capacity is [%d], now the quantity is [%d]", q.capacity, quantity)
    			return errors.New(errMsg)
    		}
    	}
    	return nil
    }
    
    /** * 获取队列中的数据 * ok: 获取成功为 true，否则false * quantity: 当前剩下的数据量，为0且ok为false则说明没有数据可读了 */
    func (q *CASQueue) getMayFail() (val interface{ }, ok bool, quantity uint32) { 
    	var putPos, getPos, getPosNew, posCnt uint32
    	var cache *casCache
    	capMod := q.capMod
    
    	putPos = atomic.LoadUint32(&q.putPos)
    	getPos = atomic.LoadUint32(&q.getPos)
    
    	if putPos >= getPos { 
    		posCnt = putPos - getPos
    	} else { 
    		posCnt = capMod + (putPos - getPos)
    	}
    
    	if posCnt < 1 { 
    		runtime.Gosched()
    		return nil, false, posCnt
    	}
    
    	getPosNew = getPos + 1
    	if !atomic.CompareAndSwapUint32(&q.getPos, getPos, getPosNew) { 
    		runtime.Gosched()
    		return nil, false, posCnt
    	}
    
    	cache = &q.cache[getPosNew&capMod]
    
    	for { 
    		getNo := atomic.LoadUint32(&cache.getNo)
    		putNo := atomic.LoadUint32(&cache.putNo)
    		if getPosNew == getNo && getNo == putNo-q.capacity { 
    			val = cache.value
    			cache.value = nil
    			atomic.AddUint32(&cache.getNo, q.capacity)
    			return val, true, posCnt - 1
    		} else { 
    			runtime.Gosched()
    		}
    	}
    }
    
    /** * 获取队列中的数据 * emptyFlag: false 则此次没获取到数据，原因是队列为空，true则获取到了数据 */
    func (q *CASQueue) Get() (val interface{ }, emptyFlag bool) { 
    	var ok bool
    	var quantity uint32
    	var v interface{ }
    	for !ok {  // 写入失败，没拿到CAS锁，则继续写入
    		v, ok, quantity = q.getMayFail()
    		// 队列为空
    		if quantity == 0 && !ok { 
    			return nil, false
    		}
    	}
    	return v, true
    }
    
    // round 到最近的2的倍数
    func minQuantity(v uint32) uint32 { 
    	v--
    	v |= v >> 1
    	v |= v >> 2
    	v |= v >> 4
    	v |= v >> 8
    	v |= v >> 16
    	v++
    	return v
    }

测试方法：

package cas_queue
    
    import (
    	"fmt"
    	"runtime"
    	"testing"
    	"time"
    )
    
    func addFuncSucc(q *CASQueue, prefix int, l int) { 
    	for i := 0; i < l; i++ { 
    		v := fmt.Sprintf("%d---%d", prefix, i)
    		err := q.Put(v)
    		if err != nil { 
    			fmt.Println(err)
    		}
    		fmt.Printf("%v ", v)
    	}
    	fmt.Println()
    }
    
    func getFuncSucc(q *CASQueue, l int) { 
    	for i := 0; i < l; i++ { 
    		_, flag := q.Get()
    		if !flag { 
    			fmt.Printf("get fail, the queue is empty\n")
    		}
    	}
    }
    
    // 测试并发时，查看是否做到了线程安全
    func TestPutManayTimesSucc(t *testing.T) { 
    	runtime.GOMAXPROCS(runtime.NumCPU())
    
    	q := NewCASQueue(10000)
    	l := 1000
    	for i := 0; i < 10; i++ { 
    		prefix := i
    		go func() { 
    			fmt.Printf("truely prefix [%d]\n", prefix)
    			addFuncSucc(q, prefix, l)
    		}()
    	}
    	time.Sleep(1 * time.Second)
    	fmt.Println("going")
    	time.Sleep(1 * time.Second)
    	fmt.Println(q.Quantity())
    
    	for i := 0; i < 10; i++ { 
    		go func() { 
    			getFuncSucc(q, l)
    		}()
    	}
    
    	time.Sleep(1 * time.Second)
    	fmt.Println("going")
    	time.Sleep(1 * time.Second)
    	fmt.Println(q.Quantity())
    }