gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/pool/swap-pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"fmt"
     5  	"gitee.com/sy_183/go-common/sgr"
     6  	"sync/atomic"
     7  )
     8  
     9  type SwapPool[O any] struct {
    10  	// queue capacity
    11  	qc int
    12  	// queue pool
    13  	qp *SyncPool[*[]O]
    14  	// consumer queue
    15  	cq atomic.Pointer[[]O]
    16  	// producer queue
    17  	pq atomic.Pointer[[]O]
    18  	// wait queue
    19  	wq atomic.Pointer[[]O]
    20  	// object creator
    21  	creator func(p Pool[O]) O
    22  }
    23  
    24  func NewSwapPool[O any](cap int, creator func(p Pool[O]) O) *SwapPool[O] {
    25  	p := &SwapPool[O]{
    26  		qc: cap,
    27  		qp: NewSyncPool(func(p *SyncPool[*[]O]) *[]O {
    28  			qp := new([]O)
    29  			*qp = make([]O, 0, cap)
    30  			return qp
    31  		}),
    32  		creator: creator,
    33  	}
    34  	p.cq.Store(p.qp.Get())
    35  	p.pq.Store(p.qp.Get())
    36  	return p
    37  }
    38  
    39  func (s *SwapPool[O]) Get() (o O) {
    40  	cqp := s.cq.Swap(nil)
    41  	fmt.Println(sgr.WrapColor("get consumer queue", sgr.FgBlue))
    42  	if cqp == nil {
    43  		panic(fmt.Errorf("swap pool does not allow concurrent get object"))
    44  	}
    45  	pq := *cqp
    46  	l := len(pq)
    47  	if l == 0 {
    48  		fmt.Println(sgr.WrapColor("get from creator start", sgr.FgRed))
    49  		o = s.creator(s)
    50  		fmt.Println(sgr.WrapColor("get from creator end", sgr.FgRed))
    51  	} else {
    52  		o = pq[l-1]
    53  		*cqp = pq[:l-1]
    54  		fmt.Println(sgr.WrapColor("get from consumer queue", sgr.FgGreen))
    55  	}
    56  	if !s.cq.CompareAndSwap(nil, cqp) {
    57  		fmt.Println(sgr.WrapColor("consumer queue swapped, put queue to wait queue", sgr.FgCyan))
    58  		if wqp := s.wq.Swap(cqp); wqp != nil {
    59  			fmt.Println(sgr.WrapColor("wait queue has last queue, put last queue to queue pool", sgr.FgMagenta))
    60  			*wqp = (*wqp)[:0]
    61  			s.qp.Put(wqp)
    62  		}
    63  	} else {
    64  		fmt.Println(sgr.WrapColor("revert consumer queue", sgr.FgBlue))
    65  	}
    66  	return
    67  }
    68  
    69  func (s *SwapPool[O]) Put(val O) {
    70  	pqp := s.pq.Load()
    71  	fmt.Println(sgr.WrapColor("get producer queue", sgr.FgBrightBlue))
    72  	if pqp == nil {
    73  		fmt.Println(sgr.WrapColor("producer queue is nil, get it from wait queue", sgr.FgBrightGreen))
    74  		pqp = s.wq.Swap(nil)
    75  		if pqp == nil {
    76  			fmt.Println(sgr.WrapColor("wait queue is nil, get producer queue from queue pool", sgr.FgBrightMagenta))
    77  			pqp = s.qp.Get()
    78  		}
    79  		s.pq.Store(pqp)
    80  	}
    81  	fmt.Println(sgr.WrapColor("put object to producer queue", sgr.FgBrightGreen))
    82  	cq := append(*pqp, val)
    83  	*pqp = cq
    84  	if len(cq) >= cap(cq) {
    85  		fmt.Println(sgr.WrapColor("producer queue full, put it to consumer queue", sgr.FgBrightBlue))
    86  		cqp := s.cq.Swap(pqp)
    87  		if cqp != nil {
    88  			fmt.Println(sgr.WrapColor("old consumer queue not used, clear and put it to producer queue", sgr.FgBrightBlue))
    89  			*cqp = (*cqp)[:0]
    90  		} else {
    91  			fmt.Println(sgr.WrapColor("old consumer used, set producer queue nil", sgr.FgBrightCyan))
    92  		}
    93  		s.pq.Store(cqp)
    94  	}
    95  }