Skip to content

Commit 2e875bc

Browse files
cyfung1031Copilot
cyfung1031
and
Copilot
authored
⚡️ 优化 MessageConnect 的记忆管理及Cleanup机制 (#1248)
* 优化 MessageConnect 的记忆管理及Cleanup机制 * handler -> cleanup * 次序可推前 * 抛出错误是合理的，不应该屏蔽错误 * Update packages/message/mock_message.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * fix --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 2c07d6b commit 2e875bc

5 files changed

Lines changed: 124 additions & 30 deletions

File tree

‎packages/message/extension_message.ts‎

Lines changed: 60 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1+
import EventEmitter from "eventemitter3";
12
import type { Message, MessageConnect, MessageSend, RuntimeMessageSender, TMessage, TMessageCommAction } from "./types";
3+
import { uuidv4 } from "@App/pkg/utils/uuid";
4+
5+
const listenerMgr = new EventEmitter<string, any>(); // 单一管理器
26

37
export class ExtensionMessage implements Message {
48
constructor(private backgroundPrimary = false) {}
@@ -19,8 +23,6 @@ export class ExtensionMessage implements Message {
1923
if (lastError) {
2024
console.error("chrome.runtime.lastError in chrome.runtime.sendMessage:", lastError);
2125
// 通信API出错不回继续对话
22-
resolve = null;
23-
return;
2426
}
2527
resolve!(resp);
2628
resolve = null;
@@ -146,25 +148,75 @@ export class ExtensionMessage implements Message {
146148
}
147149

148150
export class ExtensionMessageConnect implements MessageConnect {
149-
constructor(private con: chrome.runtime.Port) {}
151+
private readonly listenerId = `${uuidv4()}`; // 使用 uuidv4 确保唯一
152+
private con: chrome.runtime.Port | null;
153+
private isSelfDisconnected = false;
154+
155+
constructor(con: chrome.runtime.Port) {
156+
this.con = con; // 强引用
157+
const handler = (msg: TMessage, _con: chrome.runtime.Port) => {
158+
listenerMgr.emit(`onMessage:${this.listenerId}`, msg);
159+
};
160+
const cleanup = (con: chrome.runtime.Port) => {
161+
if (this.con) {
162+
this.con = null;
163+
listenerMgr.removeAllListeners(`cleanup:${this.listenerId}`);
164+
con.onMessage.removeListener(handler);
165+
con.onDisconnect.removeListener(cleanup);
166+
listenerMgr.emit(`onDisconnect:${this.listenerId}`, this.isSelfDisconnected);
167+
listenerMgr.removeAllListeners(`onDisconnect:${this.listenerId}`);
168+
listenerMgr.removeAllListeners(`onMessage:${this.listenerId}`);
169+
}
170+
};
171+
con.onMessage.addListener(handler);
172+
con.onDisconnect.addListener(cleanup);
173+
listenerMgr.once(`cleanup:${this.listenerId}`, cleanup);
174+
}
150175

151176
sendMessage(data: TMessage) {
152-
this.con.postMessage(data);
177+
if (!this.con) {
178+
console.warn("Attempted to sendMessage on a disconnected port.");
179+
// 無法 sendMessage 不应该屏蔽错误
180+
throw new Error("Attempted to sendMessage on a disconnected port.");
181+
}
182+
this.con?.postMessage(data);
153183
}
154184

155185
onMessage(callback: (data: TMessage) => void) {
156-
this.con.onMessage.addListener(callback);
186+
if (!this.con) {
187+
console.error("onMessage Invalid Port");
188+
// 無法監聽的話不应该屏蔽错误
189+
throw new Error("onMessage Invalid Port");
190+
}
191+
listenerMgr.addListener(`onMessage:${this.listenerId}`, callback);
157192
}
158193

159194
disconnect() {
160-
this.con.disconnect();
195+
if (!this.con) {
196+
console.warn("Attempted to disconnect on a disconnected port.");
197+
// 重复 disconnect() 不应该屏蔽错误
198+
throw new Error("Attempted to disconnect on a disconnected port.");
199+
}
200+
this.isSelfDisconnected = true;
201+
this.con?.disconnect();
202+
// Note: .disconnect() will NOT automatically trigger the 'cleanup' listener
203+
listenerMgr.emit(`cleanup:${this.listenerId}`);
161204
}
162205

163-
onDisconnect(callback: () => void) {
164-
this.con.onDisconnect.addListener(callback);
206+
onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
207+
if (!this.con) {
208+
console.error("onDisconnect Invalid Port");
209+
// 無法監聽的話不应该屏蔽错误
210+
throw new Error("onDisconnect Invalid Port");
211+
}
212+
listenerMgr.once(`onDisconnect:${this.listenerId}`, callback);
165213
}
166214

167215
getPort(): chrome.runtime.Port {
216+
if (!this.con) {
217+
console.error("Port is already disconnected.");
218+
throw new Error("Port is already disconnected.");
219+
}
168220
return this.con;
169221
}
170222
}

‎packages/message/mock_message.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@ export class MockMessageConnect implements MessageConnect {
1616
}
1717

1818
disconnect(): void {
19-
this.EE.emit("disconnect");
19+
this.EE.emit("disconnect", true); // MockMessageConnect 未有模拟由另一端触发 disconnect() 的情况
2020
}
2121

22-
onDisconnect(callback: () => void): void {
23-
this.EE.on("disconnect", callback);
22+
onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
23+
this.EE.once("disconnect", callback);
2424
}
2525
}
2626

‎packages/message/server.ts‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -273,15 +273,15 @@ export function forwardMessage(
273273
senderTo: MessageSend,
274274
middleware?: ApiFunctionSync
275275
) {
276-
const handler = (params: any, fromCon: IGetSender) => {
277-
const fromConnect = fromCon.getConnect();
276+
const handler = async (params: any, fromCon: IGetSender): Promise<any> => {
277+
const fromConnect: MessageConnect | undefined = fromCon.getConnect();
278278
if (fromConnect) {
279-
connect(senderTo, `${prefix}/${path}`, params).then((toCon: MessageConnect) => {
280-
fromConnect.onMessage(toCon.sendMessage.bind(toCon));
281-
toCon.onMessage(fromConnect.sendMessage.bind(fromConnect));
282-
fromConnect.onDisconnect(toCon.disconnect.bind(toCon));
283-
toCon.onDisconnect(fromConnect.disconnect.bind(fromConnect));
284-
});
279+
const toCon: MessageConnect = await connect(senderTo, `${prefix}/${path}`, params);
280+
fromConnect.onMessage(toCon.sendMessage.bind(toCon));
281+
toCon.onMessage(fromConnect.sendMessage.bind(fromConnect));
282+
fromConnect.onDisconnect(toCon.disconnect.bind(toCon));
283+
toCon.onDisconnect(fromConnect.disconnect.bind(fromConnect));
284+
return undefined;
285285
} else {
286286
return sendMessage(senderTo, prefix + "/" + path, params);
287287
}

‎packages/message/types.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ export interface MessageConnect {
5050
onMessage(callback: (data: TMessage) => void): void;
5151
sendMessage(data: TMessage): void;
5252
disconnect(): void;
53-
onDisconnect(callback: () => void): void;
53+
onDisconnect(callback: (isSelfDisconnected: boolean) => void): void;
5454
}
5555

5656
export type ExtMessageSender = {

‎packages/message/window_message.ts‎

Lines changed: 52 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ import type { Message, MessageConnect, MessageSend, RuntimeMessageSender, TMessa
22
import { uuidv4 } from "@App/pkg/utils/uuid";
33
import EventEmitter from "eventemitter3";
44

5+
const listenerMgr = new EventEmitter<string, any>(); // 单一管理器
6+
57
// 通过 window.postMessage/onmessage 实现通信
68

79
export interface PostMessage {
@@ -121,19 +123,41 @@ export class WindowMessage implements Message {
121123
}
122124

123125
export class WindowMessageConnect implements MessageConnect {
126+
private readonly listenerId = `${uuidv4()}`; // 使用 uuidv4 确保唯一
127+
private target: PostMessage | null;
128+
private isSelfDisconnected = false;
129+
124130
constructor(
125131
private messageId: string,
126-
private EE: EventEmitter<string, any>,
127-
private target: PostMessage
132+
EE: EventEmitter<string, any>,
133+
target: PostMessage
128134
) {
129-
this.onDisconnect(() => {
130-
// 移除所有监听
131-
this.EE.removeAllListeners("connectMessage:" + this.messageId);
132-
this.EE.removeAllListeners("disconnect:" + this.messageId);
133-
});
135+
this.target = target; // 强引用
136+
const handler = (msg: TMessage) => {
137+
listenerMgr.emit(`onMessage:${this.listenerId}`, msg);
138+
};
139+
const cleanup = () => {
140+
if (this.target) {
141+
this.target = null;
142+
listenerMgr.removeAllListeners(`cleanup:${this.listenerId}`);
143+
EE.removeAllListeners("connectMessage:" + this.messageId); // 模拟 con.onMessage.removeListener
144+
EE.removeAllListeners("disconnect:" + this.messageId); // 模拟 con.onDisconnect.removeListener
145+
listenerMgr.emit(`onDisconnect:${this.listenerId}`, this.isSelfDisconnected);
146+
listenerMgr.removeAllListeners(`onDisconnect:${this.listenerId}`);
147+
listenerMgr.removeAllListeners(`onMessage:${this.listenerId}`);
148+
}
149+
};
150+
EE.addListener(`connectMessage:${this.messageId}`, handler); // 模拟 con.onMessage.addListener
151+
EE.addListener(`disconnect:${this.messageId}`, cleanup); // 模拟 con.onDisconnect.addListener
152+
listenerMgr.once(`cleanup:${this.listenerId}`, cleanup);
134153
}
135154

136155
sendMessage(data: TMessage) {
156+
if (!this.target) {
157+
console.error("Attempted to sendMessage on a disconnected Target.");
158+
// 無法 sendMessage 不应该屏蔽错误
159+
throw new Error("Attempted to sendMessage on a disconnected Target.");
160+
}
137161
const body: WindowMessageBody<TMessage> = {
138162
messageId: this.messageId,
139163
type: "connectMessage",
@@ -143,20 +167,38 @@ export class WindowMessageConnect implements MessageConnect {
143167
}
144168

145169
onMessage(callback: (data: TMessage) => void) {
146-
this.EE.addListener(`connectMessage:${this.messageId}`, callback);
170+
if (!this.target) {
171+
console.error("onMessage Invalid Target");
172+
// 無法監聽的話不应该屏蔽错误
173+
throw new Error("onMessage Invalid Target");
174+
}
175+
listenerMgr.addListener(`onMessage:${this.listenerId}`, callback);
147176
}
148177

149178
disconnect() {
179+
if (!this.target) {
180+
console.warn("Attempted to disconnect on a disconnected Target.");
181+
// 重复 disconnect() 不应该屏蔽错误
182+
throw new Error("Attempted to disconnect on a disconnected Target.");
183+
}
184+
this.isSelfDisconnected = true;
150185
const body: WindowMessageBody<TMessage> = {
151186
messageId: this.messageId,
152187
type: "disconnect",
153188
data: null,
154189
};
155190
this.target.postMessage(body);
191+
// Note: .disconnect() will NOT automatically trigger the 'cleanup' listener
192+
listenerMgr.emit(`cleanup:${this.listenerId}`);
156193
}
157194

158-
onDisconnect(callback: () => void) {
159-
this.EE.addListener(`disconnect:${this.messageId}`, callback);
195+
onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
196+
if (!this.target) {
197+
console.error("onDisconnect Invalid Target");
198+
// 無法監聽的話不应该屏蔽错误
199+
throw new Error("onDisconnect Invalid Target");
200+
}
201+
listenerMgr.once(`onDisconnect:${this.listenerId}`, callback);
160202
}
161203
}
162204

0 commit comments

Comments
 (0)