增加【实时拉取消息】功能。

Author: dushixiang

pom.xml

@@ -80,6 +80,11 @@
             <artifactId>kafka-clients</artifactId>
             <version>2.8.0</version>
         </dependency>
+
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-webflux</artifactId>
+        </dependency>
     </dependencies>
 
     <build>

src/main/java/cn/typesafe/km/config/WebFluxConfig.java

@@ -0,0 +1,19 @@
+package cn.typesafe.km.config;
+
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.reactive.config.CorsRegistry;
+import org.springframework.web.reactive.config.WebFluxConfigurer;
+
+@Configuration
+public class WebFluxConfig implements WebFluxConfigurer {
+
+    @Override
+    public void addCorsMappings(CorsRegistry registry) {
+        registry.addMapping("/**")
+            .allowedOrigins("*")
+            .allowCredentials(false)
+            .allowedMethods("GET", "POST", "PUT", "DELETE", "OPTIONS")
+            .allowedHeaders("*")
+            .exposedHeaders("*");
+    }
+}

src/main/java/cn/typesafe/km/controller/TopicController.java

@@ -5,7 +5,9 @@
 import cn.typesafe.km.service.MessageService;
 import cn.typesafe.km.service.TopicService;
 import cn.typesafe.km.service.dto.*;
+import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.web.bind.annotation.*;
+import reactor.core.publisher.Flux;
 
 import javax.annotation.Resource;
 import java.util.ArrayList;
@@ -96,6 +98,14 @@ public List<ConsumerMessage> data(@PathVariable String topic, @RequestParam Stri
         return messageService.data(clusterId, topic, partition, offset, count, keyFilter, valueFilter);
     }
 
+    @GetMapping("/{topic}/data/live")
+    public Flux<ServerSentEvent<LiveMessage>> liveData(@PathVariable String topic, @RequestParam String clusterId,
+                                                                 @RequestParam(defaultValue = "0") Integer partition,
+                                                                 String keyFilter,
+                                                                 String valueFilter) {
+        return messageService.liveData(clusterId, topic, partition, keyFilter, valueFilter);
+    }
+
     @PostMapping("/{topic}/data")
     public long data(@PathVariable String topic, @RequestParam String clusterId, @RequestBody TopicData topicData) {
         return messageService.sendData(clusterId, topic, topicData);

src/main/java/cn/typesafe/km/interceptor/AuthInterceptor.java

@@ -31,6 +31,10 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
         }
 
         String tokenHeader = request.getHeader("X-Auth-Token");
+        if (!StringUtils.hasText(tokenHeader)) {
+            tokenHeader = request.getParameter("X-Auth-Token");
+        }
+
         if (StringUtils.hasText(tokenHeader)) {
             User user = tokenManager.getIfPresent(tokenHeader);
             if (user != null) {

src/main/java/cn/typesafe/km/service/MessageService.java

@@ -2,6 +2,7 @@
 
 import cn.typesafe.km.entity.Cluster;
 import cn.typesafe.km.service.dto.ConsumerMessage;
+import cn.typesafe.km.service.dto.LiveMessage;
 import cn.typesafe.km.service.dto.TopicData;
 import lombok.SneakyThrows;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -10,9 +11,11 @@
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
+import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.stereotype.Service;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
+import reactor.core.publisher.Flux;
 
 import javax.annotation.Resource;
 import java.time.Duration;
@@ -81,7 +84,81 @@ public List<ConsumerMessage> data(String clusterId, String topicName, Integer tP
             }
 
             return records
-                    .subList(0, Math.min(count, records.size()))
+                .subList(0, Math.min(count, records.size()))
+                .stream()
+                .map(record -> {
+                    int partition = record.partition();
+                    long timestamp = record.timestamp();
+                    String key = record.key();
+                    String value = record.value();
+                    long offset = record.offset();
+
+                    ConsumerMessage consumerMessage = new ConsumerMessage();
+                    consumerMessage.setTopic(topicName);
+                    consumerMessage.setOffset(offset);
+                    consumerMessage.setPartition(partition);
+                    consumerMessage.setTimestamp(timestamp);
+                    consumerMessage.setKey(key);
+                    consumerMessage.setValue(value);
+
+                    return consumerMessage;
+                }).collect(Collectors.toList());
+        }
+    }
+
+    @SneakyThrows
+    public long sendData(String clusterId, String topic, TopicData topicData) {
+        Cluster cluster = clusterService.findById(clusterId);
+        KafkaProducer<String, String> kafkaProducer = clusterService.createProducer(cluster.getServers(), cluster.getSecurityProtocol(), cluster.getSaslMechanism(), cluster.getAuthUsername(), cluster.getAuthPassword());
+        ProducerRecord<String, String> producerRecord = new ProducerRecord<>(topic, topicData.getPartition(), topicData.getKey(), topicData.getValue());
+        RecordMetadata recordMetadata = kafkaProducer.send(producerRecord).get();
+        return recordMetadata.offset();
+    }
+
+    public Flux<ServerSentEvent<LiveMessage>> liveData(String clusterId, String topicName, Integer tPartition, String keyFilter, String valueFilter) {
+
+        KafkaConsumer<String, String> kafkaConsumer = clusterService.createConsumer(clusterId);
+        TopicPartition topicPartition = new TopicPartition(topicName, tPartition);
+        List<TopicPartition> topicPartitions = Collections.singletonList(topicPartition);
+        kafkaConsumer.assign(topicPartitions);
+
+        Long endOffset = kafkaConsumer.endOffsets(topicPartitions).get(topicPartition);
+        kafkaConsumer.seek(topicPartition, endOffset);
+
+        return Flux
+            .interval(Duration.ofSeconds(1))
+            .doFinally(x -> {
+                kafkaConsumer.close();
+            })
+            .map(sequence -> {
+
+                List<ConsumerRecord<String, String>> records = new ArrayList<>();
+
+                List<ConsumerRecord<String, String>> polled = kafkaConsumer.poll(Duration.ofMillis(200)).records(topicPartition);
+
+                if (!CollectionUtils.isEmpty(polled)) {
+
+                    for (ConsumerRecord<String, String> consumerRecord : polled) {
+                        if (StringUtils.hasText(keyFilter)) {
+                            String key = consumerRecord.key();
+                            if (StringUtils.hasText(key) && key.toLowerCase().contains(keyFilter.toLowerCase())) {
+                                records.add(consumerRecord);
+                            }
+                            continue;
+                        }
+
+                        if (StringUtils.hasText(valueFilter)) {
+                            String value = consumerRecord.value();
+                            if (StringUtils.hasText(value) && value.toLowerCase().contains(valueFilter.toLowerCase())) {
+                                records.add(consumerRecord);
+                            }
+                            continue;
+                        }
+                        records.add(consumerRecord);
+                    }
+                }
+
+                List<ConsumerMessage> data = records
                     .stream()
                     .map(record -> {
                         int partition = record.partition();
@@ -100,15 +177,21 @@ public List<ConsumerMessage> data(String clusterId, String topicName, Integer tP
 
                         return consumerMessage;
                     }).collect(Collectors.toList());
-        }
-    }
 
-    @SneakyThrows
-    public long sendData(String clusterId, String topic, TopicData topicData) {
-        Cluster cluster = clusterService.findById(clusterId);
-        KafkaProducer<String, String> kafkaProducer = clusterService.createProducer(cluster.getServers(), cluster.getSecurityProtocol(), cluster.getSaslMechanism(), cluster.getAuthUsername(), cluster.getAuthPassword());
-        ProducerRecord<String, String> producerRecord = new ProducerRecord<>(topic,topicData.getPartition(), topicData.getKey(), topicData.getValue());
-        RecordMetadata recordMetadata = kafkaProducer.send(producerRecord).get();
-        return recordMetadata.offset();
+                Long currBeginningOffset = kafkaConsumer.beginningOffsets(topicPartitions).get(topicPartition);
+                Long currEndOffset = kafkaConsumer.endOffsets(topicPartitions).get(topicPartition);
+
+                LiveMessage liveMessage = new LiveMessage();
+                liveMessage.setBeginningOffset(currBeginningOffset);
+                liveMessage.setEndOffset(currEndOffset);
+                liveMessage.setPartition(tPartition);
+                liveMessage.setMessages(data);
+
+                return ServerSentEvent.<LiveMessage>builder()
+                    .id(String.valueOf(sequence))
+                    .event("topic-message-event")
+                    .data(liveMessage)
+                    .build();
+            });
     }
 }

src/main/java/cn/typesafe/km/service/dto/LiveMessage.java

@@ -0,0 +1,13 @@
+package cn.typesafe.km.service.dto;
+
+import lombok.Data;
+
+import java.util.List;
+
+@Data
+public class LiveMessage {
+    private int partition;
+    private long beginningOffset;
+    private long endOffset;
+    private List<ConsumerMessage> messages;
+}

src/main/web/src/App.js

@@ -25,6 +25,7 @@ import {
     GithubOutlined
 } from '@ant-design/icons';
 import Info from "./components/Info";
+import TopicDataLive from "./components/TopicDataLive";
 
 const {Header, Content, Footer} = Layout;
 
@@ -182,6 +183,7 @@ class App extends Component {
                                                 <Route path="/consumer-group-info" component={ConsumerGroupInfo}/>
                                                 <Route path="/topic-info" component={TopicInfo}/>
                                                 <Route path="/topic-data" component={TopicData}/>
+                                                <Route path="/topic-data-live" component={TopicDataLive}/>
                                                 <Route path="/info" component={Info}/>
                                             </Content>
                                         </Layout>

src/main/web/src/components/Topic.js

@@ -224,6 +224,12 @@ class Topic extends Component {
                                     </Button>
                                 </Link>
 
+                                <Link to={`/topic-data-live?clusterId=${this.state.clusterId}&topic=${record['name']}`}>
+                                    <Button key="2" type="link" size='small'>
+                                        <FormattedMessage id="consume-message-live" />
+                                    </Button>
+                                </Link>
+
                                 <Button type="link" size='small' onClick={() => {
                                     this.setState({
                                         createPartitionVisible: true,