23
23
import io .rsocket .resume .ResumeStateHolder ;
24
24
import java .time .Duration ;
25
25
import java .util .concurrent .TimeUnit ;
26
- import java .util .concurrent .atomic .AtomicBoolean ;
26
+ import java .util .concurrent .atomic .AtomicIntegerFieldUpdater ;
27
27
import java .util .function .Consumer ;
28
28
import reactor .core .Disposable ;
29
29
import reactor .core .publisher .Flux ;
@@ -38,11 +38,19 @@ public abstract class KeepAliveSupport implements KeepAliveFramesAcceptor {
38
38
final Duration keepAliveTimeout ;
39
39
final long keepAliveTimeoutMillis ;
40
40
41
- final AtomicBoolean started = new AtomicBoolean ();
41
+ volatile int state ;
42
+ static final AtomicIntegerFieldUpdater <KeepAliveSupport > STATE =
43
+ AtomicIntegerFieldUpdater .newUpdater (KeepAliveSupport .class , "state" );
44
+
45
+ static final int STOPPED_STATE = 0 ;
46
+ static final int STARTING_STATE = 1 ;
47
+ static final int STARTED_STATE = 2 ;
48
+ static final int DISPOSED_STATE = -1 ;
42
49
43
50
volatile Consumer <KeepAlive > onTimeout ;
44
51
volatile Consumer <ByteBuf > onFrameSent ;
45
- volatile Disposable ticksDisposable ;
52
+
53
+ Disposable ticksDisposable ;
46
54
47
55
volatile ResumeStateHolder resumeStateHolder ;
48
56
volatile long lastReceivedMillis ;
@@ -57,25 +65,30 @@ private KeepAliveSupport(
57
65
}
58
66
59
67
public KeepAliveSupport start () {
60
- this .lastReceivedMillis = scheduler .now (TimeUnit .MILLISECONDS );
61
- if (started .compareAndSet (false , true )) {
62
- ticksDisposable =
68
+ if (this .state == STOPPED_STATE && STATE .compareAndSet (this , STOPPED_STATE , STARTING_STATE )) {
69
+ this .lastReceivedMillis = scheduler .now (TimeUnit .MILLISECONDS );
70
+
71
+ final Disposable disposable =
63
72
Flux .interval (keepAliveInterval , scheduler ).subscribe (v -> onIntervalTick ());
73
+ this .ticksDisposable = disposable ;
74
+
75
+ if (this .state != STARTING_STATE
76
+ || !STATE .compareAndSet (this , STARTING_STATE , STARTED_STATE )) {
77
+ disposable .dispose ();
78
+ }
64
79
}
65
80
return this ;
66
81
}
67
82
68
83
public void stop () {
69
- if (started .compareAndSet (true , false )) {
70
- ticksDisposable .dispose ();
71
- }
84
+ terminate (STOPPED_STATE );
72
85
}
73
86
74
87
@ Override
75
88
public void receive (ByteBuf keepAliveFrame ) {
76
89
this .lastReceivedMillis = scheduler .now (TimeUnit .MILLISECONDS );
77
90
if (resumeStateHolder != null ) {
78
- long remoteLastReceivedPos = remoteLastReceivedPosition (keepAliveFrame );
91
+ final long remoteLastReceivedPos = KeepAliveFrameCodec . lastPosition (keepAliveFrame );
79
92
resumeStateHolder .onImpliedPosition (remoteLastReceivedPos );
80
93
}
81
94
if (KeepAliveFrameCodec .respondFlag (keepAliveFrame )) {
@@ -104,6 +117,16 @@ public KeepAliveSupport onTimeout(Consumer<KeepAlive> onTimeout) {
104
117
return this ;
105
118
}
106
119
120
+ @ Override
121
+ public void dispose () {
122
+ terminate (DISPOSED_STATE );
123
+ }
124
+
125
+ @ Override
126
+ public boolean isDisposed () {
127
+ return ticksDisposable .isDisposed ();
128
+ }
129
+
107
130
abstract void onIntervalTick ();
108
131
109
132
void send (ByteBuf frame ) {
@@ -122,22 +145,24 @@ void tryTimeout() {
122
145
}
123
146
}
124
147
125
- long localLastReceivedPosition ( ) {
126
- return resumeStateHolder != null ? resumeStateHolder . impliedPosition () : 0 ;
127
- }
148
+ void terminate ( int terminationState ) {
149
+ for (; ; ) {
150
+ final int state = this . state ;
128
151
129
- long remoteLastReceivedPosition ( ByteBuf keepAliveFrame ) {
130
- return KeepAliveFrameCodec . lastPosition ( keepAliveFrame ) ;
131
- }
152
+ if ( state == STOPPED_STATE || state == DISPOSED_STATE ) {
153
+ return ;
154
+ }
132
155
133
- @ Override
134
- public void dispose () {
135
- stop ();
156
+ final Disposable disposable = this .ticksDisposable ;
157
+ if (STATE .compareAndSet (this , state , terminationState )) {
158
+ disposable .dispose ();
159
+ return ;
160
+ }
161
+ }
136
162
}
137
163
138
- @ Override
139
- public boolean isDisposed () {
140
- return ticksDisposable .isDisposed ();
164
+ long localLastReceivedPosition () {
165
+ return resumeStateHolder != null ? resumeStateHolder .impliedPosition () : 0 ;
141
166
}
142
167
143
168
public static final class ClientKeepAliveSupport extends KeepAliveSupport {
0 commit comments