1717 */
1818package org .apache .atlas .kafka ;
1919
20- import kafka .server .KafkaConfig ;
21- import kafka .server .KafkaServer ;
22- import kafka .zookeeper .ZooKeeperClientException ;
20+ import kafka .testkit .KafkaClusterTestKit ;
21+ import kafka .testkit .TestKitNodes ;
2322import org .apache .atlas .ApplicationProperties ;
2423import org .apache .atlas .AtlasException ;
2524import org .apache .atlas .service .Service ;
26- import org .apache .atlas .util .CommandHandlerUtility ;
2725import org .apache .commons .configuration2 .Configuration ;
2826import org .apache .commons .configuration2 .ConfigurationConverter ;
27+ import org .apache .kafka .clients .CommonClientConfigs ;
2928import org .apache .kafka .clients .producer .ProducerConfig ;
30- import org .apache .kafka .common .KafkaException ;
31- import org .apache .kafka .common .utils .Time ;
32- import org .apache .zookeeper .server .NIOServerCnxnFactory ;
33- import org .apache .zookeeper .server .ServerCnxnFactory ;
34- import org .apache .zookeeper .server .ZooKeeperServer ;
29+ import org .apache .kafka .common .utils .Exit ;
30+ import org .apache .kafka .common .utils .Utils ;
3531import org .slf4j .Logger ;
3632import org .slf4j .LoggerFactory ;
3733import org .springframework .core .annotation .Order ;
3834import org .springframework .stereotype .Component ;
39- import scala .Option ;
4035
4136import javax .inject .Inject ;
4237
4338import java .io .File ;
44- import java .io .IOException ;
45- import java .net .BindException ;
46- import java .net .InetSocketAddress ;
47- import java .net .MalformedURLException ;
48- import java .net .URL ;
4939import java .util .Properties ;
40+ import java .util .concurrent .atomic .AtomicReference ;
5041
5142@ Component
5243@ Order (3 )
@@ -56,17 +47,18 @@ public class EmbeddedKafkaServer implements Service {
5647 public static final String PROPERTY_PREFIX = "atlas.kafka" ;
5748 public static final String PROPERTY_EMBEDDED = "atlas.notification.embedded" ;
5849
59- private static final String ATLAS_KAFKA_DATA = "data" ;
60- private static final int MAX_RETRY_TO_ACQUIRE_PORT = 3 ;
50+ private static final String ATLAS_KAFKA_DATA = "data" ;
6151
62- private final boolean isEmbedded ;
63- private final Properties properties ;
64- private KafkaServer kafkaServer ;
65- private ServerCnxnFactory factory ;
52+ private final boolean isEmbedded ;
53+ private final Configuration applicationProperties ;
54+ private final Properties properties ;
55+
56+ private KafkaClusterTestKit cluster ;
6657
6758 @ Inject
6859 public EmbeddedKafkaServer (Configuration applicationProperties ) {
69- Configuration kafkaConf = ApplicationProperties .getSubsetConfiguration (applicationProperties , PROPERTY_PREFIX );
60+ this .applicationProperties = applicationProperties ;
61+ Configuration kafkaConf = ApplicationProperties .getSubsetConfiguration (applicationProperties , PROPERTY_PREFIX );
7062
7163 this .isEmbedded = applicationProperties .getBoolean (PROPERTY_EMBEDDED , false );
7264 this .properties = ConfigurationConverter .getProperties (kafkaConf );
@@ -78,8 +70,7 @@ public void start() throws AtlasException {
7870
7971 if (isEmbedded ) {
8072 try {
81- startZk ();
82- startKafka ();
73+ startKraftBroker ();
8374 } catch (Exception e ) {
8475 throw new AtlasException ("Failed to start embedded kafka" , e );
8576 }
@@ -94,90 +85,77 @@ public void start() throws AtlasException {
9485 public void stop () {
9586 LOG .info ("==> EmbeddedKafkaServer.stop(isEmbedded={})" , isEmbedded );
9687
97- if (kafkaServer != null ) {
98- kafkaServer .shutdown ();
99- }
88+ if (cluster != null ) {
89+ AtomicReference <Throwable > shutdownFailure = new AtomicReference <>();
90+
91+ Utils .closeQuietly (cluster , "embedded Kafka cluster" , shutdownFailure );
92+
93+ if (shutdownFailure .get () != null ) {
94+ LOG .warn ("Failed to shut down embedded Kafka cluster" , shutdownFailure .get ());
95+ }
10096
101- if (factory != null ) {
102- factory .shutdown ();
97+ cluster = null ;
10398 }
10499
105100 LOG .info ("<== EmbeddedKafka.stop(isEmbedded={})" , isEmbedded );
106101 }
107102
108- private String startZk () throws IOException , InterruptedException {
109- String zkValue = properties . getProperty ( "zookeeper.connect" );
103+ private void startKraftBroker () throws Exception {
104+ overrideExitMethods ( );
110105
111- LOG . info ( "Starting zookeeper at {}" , zkValue );
106+ File logDir = constructDir ( "kafka" );
112107
113- URL zkAddress = getURL (zkValue );
114- File snapshotDir = constructDir ("zk/txn" );
115- File logDir = constructDir ("zk/snap" );
108+ LOG .info ("Starting embedded KRaft kafka (log.dir={})" , logDir .getAbsolutePath ());
116109
117- for ( int attemptCount = 0 ; attemptCount < MAX_RETRY_TO_ACQUIRE_PORT ; attemptCount ++) {
118- try {
119- factory = NIOServerCnxnFactory . createFactory ( new InetSocketAddress ( zkAddress . getHost (), zkAddress . getPort ()), 1024 );
120- break ;
121- } catch ( BindException e ) {
122- LOG . warn ( "Attempt {}: Starting zookeeper at {} failed" , attemptCount , zkValue );
110+ KafkaClusterTestKit . Builder clusterBuilder = new KafkaClusterTestKit . Builder (
111+ new TestKitNodes . Builder ()
112+ . setCombined ( true )
113+ . setNumBrokerNodes ( 1 )
114+ . setNumControllerNodes ( 1 )
115+ . build () );
123116
124- if (attemptCount == MAX_RETRY_TO_ACQUIRE_PORT - 1 ) {
125- throw e ;
126- }
117+ Properties brokerConfig = buildBrokerConfig (logDir );
127118
128- CommandHandlerUtility .tryKillingProcessUsingPort (zkAddress .getPort (), attemptCount != 0 );
129- }
130- }
119+ brokerConfig .forEach ((key , value ) -> clusterBuilder .setConfigProp (key .toString (), value ));
131120
132- factory .startup (new ZooKeeperServer (snapshotDir , logDir , 500 ));
121+ cluster = clusterBuilder .build ();
122+ cluster .format ();
123+ cluster .startup ();
124+ cluster .waitForReadyBrokers ();
133125
134- String ret = factory . getLocalAddress ().getAddress ( ).toString ();
126+ String bootstrapServers = cluster . clientProperties ().get ( CommonClientConfigs . BOOTSTRAP_SERVERS_CONFIG ).toString ();
135127
136- LOG .info ("Embedded zookeeper for Kafka started at {}" , ret );
128+ LOG .info ("Embedded KRaft kafka server started at {}" , bootstrapServers );
137129
138- return ret ;
130+ applicationProperties .setProperty (PROPERTY_PREFIX + ".bootstrap.servers" , bootstrapServers );
131+ properties .setProperty (ProducerConfig .BOOTSTRAP_SERVERS_CONFIG , bootstrapServers );
139132 }
140133
141- private void startKafka () throws IOException {
142- String kafkaValue = properties . getProperty ( ProducerConfig . BOOTSTRAP_SERVERS_CONFIG );
134+ private Properties buildBrokerConfig ( File logDir ) {
135+ Properties brokerConfig = new Properties ( );
143136
144- LOG .info ("Starting kafka at {}" , kafkaValue );
137+ brokerConfig .setProperty ("log.dir" , logDir .getAbsolutePath ());
138+ brokerConfig .setProperty ("delete.topic.enable" , "true" );
139+ brokerConfig .setProperty ("group.initial.rebalance.delay.ms" , "0" );
140+ brokerConfig .setProperty ("offsets.topic.replication.factor" , "1" );
141+ brokerConfig .setProperty ("transaction.state.log.replication.factor" , "1" );
142+ brokerConfig .setProperty ("transaction.state.log.min.isr" , "1" );
143+ brokerConfig .setProperty ("num.partitions" , "1" );
145144
146- URL kafkaAddress = getURL (kafkaValue );
147- Properties brokerConfig = properties ;
145+ String replicationFactor = properties .getProperty ("offsets.topic.replication.factor" );
148146
149- for (int attemptCount = 0 ; attemptCount < MAX_RETRY_TO_ACQUIRE_PORT ; attemptCount ++) {
150- try {
151- brokerConfig .setProperty ("broker.id" , "1" );
152- brokerConfig .setProperty ("host.name" , kafkaAddress .getHost ());
153- brokerConfig .setProperty ("port" , String .valueOf (kafkaAddress .getPort ()));
154- brokerConfig .setProperty ("log.dirs" , constructDir ("kafka" ).getAbsolutePath ());
155- brokerConfig .setProperty ("log.flush.interval.messages" , String .valueOf (1 ));
156-
157- kafkaServer = new KafkaServer (KafkaConfig .fromProps (brokerConfig ), Time .SYSTEM , Option .apply (this .getClass ().getName ()), false );
158-
159- kafkaServer .startup ();
160- break ;
161- } catch (KafkaException | ZooKeeperClientException e ) {
162- LOG .warn ("Attempt {}: kafka server with broker config {} failed" , attemptCount , brokerConfig );
163-
164- if (attemptCount == MAX_RETRY_TO_ACQUIRE_PORT - 1 ) {
165- throw e ;
166- }
167-
168- if (kafkaServer != null ) {
169- try {
170- kafkaServer .shutdown ();
171- } catch (Exception ex ) {
172- LOG .info ("Failed to shutdown kafka server" , ex );
173- }
174- }
175-
176- CommandHandlerUtility .tryKillingProcessUsingPort (kafkaAddress .getPort (), attemptCount != 0 );
177- }
147+ if (replicationFactor != null ) {
148+ brokerConfig .setProperty ("offsets.topic.replication.factor" , replicationFactor );
178149 }
179150
180- LOG .info ("Embedded kafka server started with broker config {}" , brokerConfig );
151+ return brokerConfig ;
152+ }
153+
154+ private void overrideExitMethods () {
155+ Exit .setExitProcedure ((statusCode , message ) ->
156+ LOG .warn ("Kafka Exit.exit({}, {}) suppressed in embedded broker" , statusCode , message ));
157+ Exit .setHaltProcedure ((statusCode , message ) ->
158+ LOG .warn ("Kafka Exit.halt({}, {}) suppressed in embedded broker" , statusCode , message ));
181159 }
182160
183161 private File constructDir (String dirPrefix ) {
@@ -189,12 +167,4 @@ private File constructDir(String dirPrefix) {
189167
190168 return file ;
191169 }
192-
193- private URL getURL (String url ) throws MalformedURLException {
194- try {
195- return new URL (url );
196- } catch (MalformedURLException e ) {
197- return new URL ("http://" + url );
198- }
199- }
200170}
0 commit comments