4040import org .apache .iotdb .consensus .ratis .utils .Utils ;
4141
4242import org .apache .ratis .thirdparty .com .google .common .base .Preconditions ;
43+ import org .apache .ratis .util .ExitUtils ;
4344import org .apache .ratis .util .FileUtils ;
4445import org .apache .ratis .util .TimeDuration ;
4546import org .apache .ratis .util .Timestamp ;
5556import java .nio .ByteBuffer ;
5657import java .util .ArrayList ;
5758import java .util .Collections ;
59+ import java .util .HashSet ;
5860import java .util .List ;
5961import java .util .Objects ;
6062import java .util .Scanner ;
63+ import java .util .Set ;
6164import java .util .concurrent .CompletableFuture ;
6265import java .util .concurrent .CountDownLatch ;
6366import java .util .concurrent .ExecutorService ;
7073public class TestUtils {
7174
7275 private static final Logger logger = LoggerFactory .getLogger (TestUtils .class );
76+ private static final TimeDuration PORT_RELEASE_WAIT = TimeDuration .valueOf (10 , TimeUnit .SECONDS );
77+ private static final TimeDuration PORT_RELEASE_POLL =
78+ TimeDuration .valueOf (100 , TimeUnit .MILLISECONDS );
7379
7480 public static class TestDataSet implements DataSet {
7581
@@ -230,6 +236,7 @@ static class MiniCluster {
230236 private final RatisConfig config ;
231237 private final List <RatisConsensus > servers ;
232238 private final ConsensusGroup group ;
239+ private final List <Integer > peerPorts ;
233240 private Supplier <IStateMachine > smProvider ;
234241 private final AtomicBoolean isStopped = new AtomicBoolean (false );
235242
@@ -250,9 +257,10 @@ private MiniCluster(
250257 this .peerStorage = new ArrayList <>();
251258 this .stateMachines = new ArrayList <>();
252259 this .servers = new ArrayList <>();
260+ this .peerPorts = randomDistinctPorts (replicas );
253261
254262 for (int i = 0 ; i < replicas ; i ++) {
255- peers .add (new Peer (gid , i , new TEndPoint ("127.0.0.1" , randomFreePort ( ))));
263+ peers .add (new Peer (gid , i , new TEndPoint ("127.0.0.1" , peerPorts . get ( i ))));
256264
257265 final File storage = storageProvider .apply (i );
258266 FileUtils .deleteFileQuietly (storage );
@@ -300,6 +308,7 @@ void stop() throws IOException {
300308 for (RatisConsensus server : servers ) {
301309 server .stop ();
302310 }
311+ waitUntilPortsReleased (peerPorts );
303312 isStopped .set (true );
304313 }
305314
@@ -495,6 +504,65 @@ MiniCluster create() {
495504 }
496505 }
497506
507+ static void prepareJvmForRatisTest () {
508+ ExitUtils .disableSystemExit ();
509+ ExitUtils .clear ();
510+ }
511+
512+ static void assertNoUnexpectedRatisExit () {
513+ ExitUtils .assertNotTerminated ();
514+ ExitUtils .clear ();
515+ }
516+
517+ private static List <Integer > randomDistinctPorts (int count ) {
518+ final List <Integer > ports = new ArrayList <>(count );
519+ final Set <Integer > uniquePorts = new HashSet <>();
520+ while (ports .size () < count ) {
521+ final int port = randomFreePort ();
522+ if (uniquePorts .add (port )) {
523+ ports .add (port );
524+ }
525+ }
526+ return ports ;
527+ }
528+
529+ private static void waitUntilPortsReleased (List <Integer > ports ) throws IOException {
530+ final Timestamp start = Timestamp .currentTime ();
531+ while (true ) {
532+ boolean allReleased = true ;
533+ for (int port : ports ) {
534+ if (!isLocalPortAvailable (port )) {
535+ allReleased = false ;
536+ break ;
537+ }
538+ }
539+
540+ if (allReleased ) {
541+ return ;
542+ }
543+
544+ if (start .elapsedTime ().compareTo (PORT_RELEASE_WAIT ) > 0 ) {
545+ throw new IOException ("Timed out waiting for Ratis test ports to be released: " + ports );
546+ }
547+
548+ try {
549+ PORT_RELEASE_POLL .sleep ();
550+ } catch (InterruptedException e ) {
551+ Thread .currentThread ().interrupt ();
552+ throw new IOException ("Interrupted while waiting for Ratis test ports to be released" , e );
553+ }
554+ }
555+ }
556+
557+ private static boolean isLocalPortAvailable (int port ) {
558+ try (ServerSocket socket = new ServerSocket (port )) {
559+ socket .setReuseAddress (true );
560+ return true ;
561+ } catch (IOException e ) {
562+ return false ;
563+ }
564+ }
565+
498566 private static int randomFreePort () {
499567 try (ServerSocket socket = new ServerSocket (0 )) {
500568 return socket .getLocalPort ();
0 commit comments