github.com/bruceshao/lockfree@v1.1.3-0.20230816090528-e89824c0a6e9/lockfree.go (about)

     1  /*
     2   * Copyright (C) THL A29 Limited, a Tencent company. All rights reserved.
     3   *
     4   * SPDX-License-Identifier: Apache-2.0
     5   *
     6   */
     7  
     8  package lockfree
     9  
    10  import (
    11  	"fmt"
    12  	"sync/atomic"
    13  )
    14  
    15  // Lockfree 包装类，内部包装了生产者和消费者
    16  type Lockfree[T any] struct {
    17  	writer   *Producer[T]
    18  	consumer *consumer[T]
    19  	status   int32
    20  }
    21  
    22  // NewLockfree 自定义创建消费端的Disruptor
    23  // capacity：buffer的容量大小，类似于chan的大小，但要求必须是2^n，即2的指数倍，如果不是的话会被修改
    24  // handler：消费端的事件处理器
    25  // blocks：读取阻塞时的处理策略
    26  func NewLockfree[T any](capacity int, handler EventHandler[T], blocks blockStrategy) *Lockfree[T] {
    27  	// 重新计算正确的容量
    28  	capacity = minSuitableCap(capacity)
    29  	seqer := newSequencer(capacity)
    30  	rbuf := newRingBuffer[T](capacity)
    31  	cmer := newConsumer[T](rbuf, handler, seqer, blocks)
    32  	writer := newProducer[T](seqer, rbuf, blocks)
    33  	return &Lockfree[T]{
    34  		writer:   writer,
    35  		consumer: cmer,
    36  		status:   READY,
    37  	}
    38  }
    39  
    40  func (d *Lockfree[T]) Start() error {
    41  	if atomic.CompareAndSwapInt32(&d.status, READY, RUNNING) {
    42  		// 启动消费者
    43  		if err := d.consumer.start(); err != nil {
    44  			// 恢复现场
    45  			atomic.CompareAndSwapInt32(&d.status, RUNNING, READY)
    46  			return err
    47  		}
    48  		// 启动生产者
    49  		if err := d.writer.start(); err != nil {
    50  			// 恢复现场
    51  			atomic.CompareAndSwapInt32(&d.status, RUNNING, READY)
    52  			return err
    53  		}
    54  		return nil
    55  	}
    56  	return fmt.Errorf(StartErrorFormat, "Disruptor")
    57  }
    58  
    59  func (d *Lockfree[T]) Producer() *Producer[T] {
    60  	return d.writer
    61  }
    62  
    63  func (d *Lockfree[T]) Running() bool {
    64  	return d.status == RUNNING
    65  }
    66  
    67  func (d *Lockfree[T]) Close() error {
    68  	if atomic.CompareAndSwapInt32(&d.status, RUNNING, READY) {
    69  		// 关闭生产者
    70  		if err := d.writer.close(); err != nil {
    71  			// 恢复现场
    72  			atomic.CompareAndSwapInt32(&d.status, READY, RUNNING)
    73  			return err
    74  		}
    75  		// 关闭消费者
    76  		if err := d.consumer.close(); err != nil {
    77  			// 恢复现场
    78  			atomic.CompareAndSwapInt32(&d.status, READY, RUNNING)
    79  			return err
    80  		}
    81  		// 关闭成功
    82  		return nil
    83  	}
    84  	return fmt.Errorf(CloseErrorFormat, "Disruptor")
    85  }