github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/mq/mq.go (about)

     1  package mq
     2  
     3  import (
     4  	"encoding/gob"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
     6  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
     7  	"strconv"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  // Message 消息事件正文
    13  type Message struct {
    14  	// 消息触发者
    15  	TriggeredBy string
    16  
    17  	// 事件标识
    18  	Event string
    19  
    20  	// 消息正文
    21  	Content interface{}
    22  }
    23  
    24  type CallbackFunc func(Message)
    25  
    26  // MQ 消息队列
    27  type MQ interface {
    28  	rpc.Notifier
    29  
    30  	// 发布一个消息
    31  	Publish(string, Message)
    32  
    33  	// 订阅一个消息主题
    34  	Subscribe(string, int) <-chan Message
    35  
    36  	// 订阅一个消息主题，注册触发回调函数
    37  	SubscribeCallback(string, CallbackFunc)
    38  
    39  	// 取消订阅一个消息主题
    40  	Unsubscribe(string, <-chan Message)
    41  }
    42  
    43  var GlobalMQ = NewMQ()
    44  
    45  func NewMQ() MQ {
    46  	return &inMemoryMQ{
    47  		topics:    make(map[string][]chan Message),
    48  		callbacks: make(map[string][]CallbackFunc),
    49  	}
    50  }
    51  
    52  func init() {
    53  	gob.Register(Message{})
    54  	gob.Register([]rpc.Event{})
    55  }
    56  
    57  type inMemoryMQ struct {
    58  	topics    map[string][]chan Message
    59  	callbacks map[string][]CallbackFunc
    60  	sync.RWMutex
    61  }
    62  
    63  func (i *inMemoryMQ) Publish(topic string, message Message) {
    64  	i.RLock()
    65  	subscribersChan, okChan := i.topics[topic]
    66  	subscribersCallback, okCallback := i.callbacks[topic]
    67  	i.RUnlock()
    68  
    69  	if okChan {
    70  		go func(subscribersChan []chan Message) {
    71  			for i := 0; i < len(subscribersChan); i++ {
    72  				select {
    73  				case subscribersChan[i] <- message:
    74  				case <-time.After(time.Millisecond * 500):
    75  				}
    76  			}
    77  		}(subscribersChan)
    78  
    79  	}
    80  
    81  	if okCallback {
    82  		for i := 0; i < len(subscribersCallback); i++ {
    83  			go subscribersCallback[i](message)
    84  		}
    85  	}
    86  }
    87  
    88  func (i *inMemoryMQ) Subscribe(topic string, buffer int) <-chan Message {
    89  	ch := make(chan Message, buffer)
    90  	i.Lock()
    91  	i.topics[topic] = append(i.topics[topic], ch)
    92  	i.Unlock()
    93  	return ch
    94  }
    95  
    96  func (i *inMemoryMQ) SubscribeCallback(topic string, callbackFunc CallbackFunc) {
    97  	i.Lock()
    98  	i.callbacks[topic] = append(i.callbacks[topic], callbackFunc)
    99  	i.Unlock()
   100  }
   101  
   102  func (i *inMemoryMQ) Unsubscribe(topic string, sub <-chan Message) {
   103  	i.Lock()
   104  	defer i.Unlock()
   105  
   106  	subscribers, ok := i.topics[topic]
   107  	if !ok {
   108  		return
   109  	}
   110  
   111  	var newSubs []chan Message
   112  	for _, subscriber := range subscribers {
   113  		if subscriber == sub {
   114  			continue
   115  		}
   116  		newSubs = append(newSubs, subscriber)
   117  	}
   118  
   119  	i.topics[topic] = newSubs
   120  }
   121  
   122  func (i *inMemoryMQ) Aria2Notify(events []rpc.Event, status int) {
   123  	for _, event := range events {
   124  		i.Publish(event.Gid, Message{
   125  			TriggeredBy: event.Gid,
   126  			Event:       strconv.FormatInt(int64(status), 10),
   127  			Content:     events,
   128  		})
   129  	}
   130  }
   131  
   132  // OnDownloadStart 下载开始
   133  func (i *inMemoryMQ) OnDownloadStart(events []rpc.Event) {
   134  	i.Aria2Notify(events, common.Downloading)
   135  }
   136  
   137  // OnDownloadPause 下载暂停
   138  func (i *inMemoryMQ) OnDownloadPause(events []rpc.Event) {
   139  	i.Aria2Notify(events, common.Paused)
   140  }
   141  
   142  // OnDownloadStop 下载停止
   143  func (i *inMemoryMQ) OnDownloadStop(events []rpc.Event) {
   144  	i.Aria2Notify(events, common.Canceled)
   145  }
   146  
   147  // OnDownloadComplete 下载完成
   148  func (i *inMemoryMQ) OnDownloadComplete(events []rpc.Event) {
   149  	i.Aria2Notify(events, common.Complete)
   150  }
   151  
   152  // OnDownloadError 下载出错
   153  func (i *inMemoryMQ) OnDownloadError(events []rpc.Event) {
   154  	i.Aria2Notify(events, common.Error)
   155  }
   156  
   157  // OnBtDownloadComplete BT下载完成
   158  func (i *inMemoryMQ) OnBtDownloadComplete(events []rpc.Event) {
   159  	i.Aria2Notify(events, common.Complete)
   160  }