1
1
import * as cbor from "cbor-x" ;
2
2
import invariant from "invariant" ;
3
3
import onChange from "on-change" ;
4
+ import { success } from "zod/v4" ;
4
5
import type { ActorKey } from "@/actor/mod" ;
5
6
import type * as wsToClient from "@/actor/protocol/message/to-client" ;
6
7
import type * as wsToServer from "@/actor/protocol/message/to-server" ;
@@ -9,10 +10,10 @@ import type { Logger } from "@/common/log";
9
10
import { isCborSerializable , stringifyError } from "@/common/utils" ;
10
11
import type { UniversalWebSocket } from "@/common/websocket-interface" ;
11
12
import { ActorInspector } from "@/inspector/actor" ;
12
- import type { Registry , RegistryConfig } from "@/mod" ;
13
+ import type { Registry } from "@/mod" ;
13
14
import type { ActionContext } from "./action" ;
14
15
import type { ActorConfig } from "./config" ;
15
- import { Conn , type ConnId } from "./connection" ;
16
+ import { Conn , type ConnId , generatePing } from "./connection" ;
16
17
import { ActorContext } from "./context" ;
17
18
import type { AnyDatabaseProvider , InferDatabaseClient } from "./database" ;
18
19
import type { ActorDriver , ConnDriver , ConnDrivers } from "./driver" ;
@@ -157,6 +158,11 @@ export class ActorInstance<
157
158
#ready = false ;
158
159
159
160
#connections = new Map < ConnId , Conn < S , CP , CS , V , I , AD , DB > > ( ) ;
161
+ // This is used to track the last ping sent to the client, when restoring a connection
162
+ #connectionsPingRequests = new Map <
163
+ ConnId ,
164
+ { ping : string ; resolve : ( ) => void }
165
+ > ( ) ;
160
166
#subscriptionIndex = new Map < string , Set < Conn < S , CP , CS , V , I , AD , DB > > > ( ) ;
161
167
162
168
#schedule! : Schedule ;
@@ -591,6 +597,8 @@ export class ActorInstance<
591
597
// Set initial state
592
598
this . #setPersist( persistData ) ;
593
599
600
+ const restorePromises = [ ] ;
601
+
594
602
// Load connections
595
603
for ( const connPersist of this . #persist. c ) {
596
604
// Create connections
@@ -601,13 +609,60 @@ export class ActorInstance<
601
609
driver ,
602
610
this . #connStateEnabled,
603
611
) ;
604
- this . #connections. set ( conn . id , conn ) ;
605
612
606
- // Register event subscriptions
607
- for ( const sub of connPersist . su ) {
608
- this . #addSubscription( sub . n , conn , true ) ;
609
- }
613
+ // send ping, to ensure the connection is alive
614
+
615
+ const { promise, resolve } = Promise . withResolvers < void > ( ) ;
616
+ restorePromises . push (
617
+ Promise . race ( [
618
+ promise ,
619
+ new Promise < void > ( ( _ , reject ) => {
620
+ setTimeout ( ( ) => {
621
+ reject ( ) ;
622
+ } , 2500 ) ;
623
+ } ) ,
624
+ ] )
625
+ . catch ( ( ) => {
626
+ process . nextTick ( ( ) => {
627
+ logger ( ) . debug ( "connection restore failed" , {
628
+ connId : conn . id ,
629
+ } ) ;
630
+ this . #connections. delete ( conn . id ) ;
631
+ conn . disconnect (
632
+ "Connection restore failed, connection is not alive" ,
633
+ ) ;
634
+ this . __removeConn ( conn ) ;
635
+ } ) ;
636
+ } )
637
+ . then ( ( ) => {
638
+ logger ( ) . debug ( "connection restored" , {
639
+ connId : conn . id ,
640
+ } ) ;
641
+ this . #connections. set ( conn . id , conn ) ;
642
+
643
+ // Register event subscriptions
644
+ for ( const sub of connPersist . su ) {
645
+ this . #addSubscription( sub . n , conn , true ) ;
646
+ }
647
+ } ) ,
648
+ ) ;
649
+
650
+ const ping = generatePing ( ) ;
651
+ this . #connectionsPingRequests. set ( conn . id , { ping, resolve } ) ;
652
+ conn . _sendMessage (
653
+ new CachedSerializer < wsToClient . ToClient > ( {
654
+ b : {
655
+ p : ping ,
656
+ } ,
657
+ } ) ,
658
+ ) ;
610
659
}
660
+
661
+ const result = await Promise . allSettled ( restorePromises ) ;
662
+ logger ( ) . info ( "connections restored" , {
663
+ success : result . filter ( ( r ) => r . status === "fulfilled" ) . length ,
664
+ failed : result . filter ( ( r ) => r . status === "rejected" ) . length ,
665
+ } ) ;
611
666
} else {
612
667
logger ( ) . info ( "actor creating" ) ;
613
668
@@ -818,6 +873,8 @@ export class ActorInstance<
818
873
this . #persist. c . push ( persist ) ;
819
874
this . saveState ( { immediate : true } ) ;
820
875
876
+ this . inspector . emitter . emit ( "connectionUpdated" ) ;
877
+
821
878
// Handle connection
822
879
if ( this . #config. onConnect ) {
823
880
try {
@@ -841,8 +898,6 @@ export class ActorInstance<
841
898
}
842
899
}
843
900
844
- this . inspector . emitter . emit ( "connectionUpdated" ) ;
845
-
846
901
// Send init message
847
902
conn . _sendMessage (
848
903
new CachedSerializer < wsToClient . ToClient > ( {
@@ -890,6 +945,14 @@ export class ActorInstance<
890
945
} ) ;
891
946
this . #removeSubscription( eventName , conn , false ) ;
892
947
} ,
948
+ onPong : async ( pong , conn ) => {
949
+ const pingRequest = this . #connectionsPingRequests. get ( conn . id ) ;
950
+ if ( pingRequest ?. ping === pong ) {
951
+ // Resolve the ping request if it matches the last sent ping
952
+ pingRequest . resolve ( ) ;
953
+ this . #connectionsPingRequests. delete ( conn . id ) ;
954
+ }
955
+ } ,
893
956
} ) ;
894
957
}
895
958
0 commit comments