返回 AiToEarn
pub-sub.service.ts
根目录 / project / aitoearn-backend / libs / redis / src / pub-sub.service.ts
1 /* eslint-disable ts/no-explicit-any */
2 import { Injectable, Logger, OnModuleInit } from '@nestjs/common'
3 import { Redis } from 'ioredis'
4
5 @Injectable()
6 export class RedisPubSubService implements OnModuleInit {
7 private readonly logger = new Logger(RedisPubSubService.name)
8 private readonly listeners = new Map<string, Set<(data: any) => void>>()
9 private readonly onceListeners = new Map<string, Set<(data: any) => void>>()
10 private readonly subscribedChannels = new Set<string>()
11
12 private readonly messageHandler = (channel: string, message: string) => {
13 try {
14 const data = JSON.parse(message)
15 this.dispatchMessage(channel, data)
16 }
17 catch (error) {
18 this.logger.error(`Failed to parse message on channel ${channel}: ${error}`)
19 }
20 }
21
22 constructor(
23 private readonly subscriber: Redis,
24 private readonly publisher: Redis,
25 ) { }
26
27 async onModuleInit() {
28 this.subscriber.on('message', this.messageHandler)
29 this.logger.log('Redis Pub/Sub EventEmitter initialized')
30 }
31
32 on<T = any>(channel: string, listener: (data: T) => void): this {
33 if (!this.listeners.has(channel)) {
34 this.listeners.set(channel, new Set())
35 }
36 this.listeners.get(channel)!.add(listener)
37
38 if (!this.subscribedChannels.has(channel)) {
39 this.subscribeChannel(channel)
40 }
41
42 return this
43 }
44
45 off<T = any>(channel: string, listener: (data: T) => void): this {
46 const channelListeners = this.listeners.get(channel)
47 if (channelListeners) {
48 channelListeners.delete(listener)
49 if (channelListeners.size === 0) {
50 this.listeners.delete(channel)
51 }
52 }
53
54 const channelOnceListeners = this.onceListeners.get(channel)
55 if (channelOnceListeners) {
56 channelOnceListeners.delete(listener)
57 if (channelOnceListeners.size === 0) {
58 this.onceListeners.delete(channel)
59 }
60 }
61
62 if (!this.listeners.has(channel) && !this.onceListeners.has(channel)) {
63 this.unsubscribeChannel(channel)
64 }
65
66 return this
67 }
68
69 once<T = any>(channel: string, listener: (data: T) => void): this {
70 const wrappedListener = (data: any) => {
71 this.off(channel, wrappedListener)
72 listener(data)
73 }
74
75 if (!this.onceListeners.has(channel)) {
76 this.onceListeners.set(channel, new Set())
77 }
78 this.onceListeners.get(channel)!.add(wrappedListener)
79
80 if (!this.subscribedChannels.has(channel)) {
81 this.subscribeChannel(channel)
82 }
83
84 return this
85 }
86
87 async emit(channel: string, data: any): Promise<boolean> {
88 try {
89 const res = await this.publisher.publish(channel, JSON.stringify(data))
90 return !!res
91 }
92 catch (error) {
93 this.logger.error(`Failed to emit message on channel ${channel}: ${error}`)
94 return false
95 }
96 }
97
98 private dispatchMessage(channel: string, data: any) {
99 const channelListeners = this.listeners.get(channel)
100 if (channelListeners) {
101 for (const listener of channelListeners) {
102 try {
103 listener(data)
104 }
105 catch (error) {
106 this.logger.error(`Error in listener for channel ${channel}: ${error}`)
107 }
108 }
109 }
110
111 const channelOnceListeners = this.onceListeners.get(channel)
112 if (channelOnceListeners) {
113 const listenersToCall = Array.from(channelOnceListeners)
114 for (const listener of listenersToCall) {
115 try {
116 listener(data)
117 }
118 catch (error) {
119 this.logger.error(`Error in once listener for channel ${channel}: ${error}`)
120 }
121 }
122 }
123 }
124
125 private subscribeChannel(channel: string) {
126 if (this.subscribedChannels.has(channel)) {
127 return
128 }
129
130 this.subscriber.subscribe(channel).catch((error) => {
131 this.logger.error(`Failed to subscribe to channel ${channel}: ${error}`)
132 })
133 this.subscribedChannels.add(channel)
134 this.logger.debug(`Auto-subscribed to channel: ${channel}`)
135 }
136
137 private unsubscribeChannel(channel: string) {
138 if (!this.subscribedChannels.has(channel)) {
139 return
140 }
141
142 this.subscriber.unsubscribe(channel).catch((error) => {
143 this.logger.error(`Failed to unsubscribe from channel ${channel}: ${error}`)
144 })
145 this.subscribedChannels.delete(channel)
146 this.logger.debug(`Auto-unsubscribed from channel: ${channel}`)
147 }
148 }
149
149 lines TYPESCRIPT