crawler-plugin/backend-java/src/test/java/com/nanri/aiimage/config/HttpClientConnectionReuseTest2.java

package com.nanri.aiimage.config;

import org.junit.jupiter.api.Test;
import org.springframework.web.client.RestClient;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.http.HttpClient;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * task-174：连接复用测试（plan 10）。
 *
 * 用本机计数 HTTP/1.1 服务统计真实 TCP 连接数，验证 HttpClientPool 共享池的 keep-alive
 * 复用：连续顺序请求复用同一连接、不同客户端工厂共用同一连接池、同 host 突发不新建连接、
 * 不同 host 各自建连、非共享池客户端不复用。全部本地确定性，不发外部网络请求。
 *
 * 平台边界：JDK HttpClient HTTP/1.1 无多路复用，并发在途请求会各自建连；复用体现在
 * 顺序/间歇请求复用空闲连接上，故并发后追加的顺序请求应复用已建立的连接而不新增。
 */
class HttpClientConnectionReuseTest2 {

    /** 极小的本地 HTTP/1.1 keep-alive 计数服务。 */
    static final class CountingServer implements AutoCloseable {
        private final ServerSocket serverSocket;
        private final AtomicInteger accepted = new AtomicInteger();
        private final AtomicInteger requests = new AtomicInteger();
        private final List<Socket> sockets = new CopyOnWriteArrayList<>();
        private final Thread acceptThread;
        private volatile boolean running = true;

        CountingServer() throws IOException {
            serverSocket = new ServerSocket(0, 256, InetAddress.getByName("127.0.0.1"));
            acceptThread = new Thread(this::acceptLoop, "conn-reuse-srv");
            acceptThread.setDaemon(true);
            acceptThread.start();
        }

        int port() {
            return serverSocket.getLocalPort();
        }

        int accepted() {
            return accepted.get();
        }

        int requests() {
            return requests.get();
        }

        private void acceptLoop() {
            while (running) {
                try {
                    Socket socket = serverSocket.accept();
                    sockets.add(socket);
                    accepted.incrementAndGet();
                    Thread worker = new Thread(() -> handle(socket), "conn-reuse-worker");
                    worker.setDaemon(true);
                    worker.start();
                } catch (IOException e) {
                    if (running) {
                        // accept 被关闭时退出
                    }
                    return;
                }
            }
        }

        private void handle(Socket socket) {
            try (socket) {
                InputStream in = socket.getInputStream();
                OutputStream out = socket.getOutputStream();
                while (running && readRequest(in)) {
                    requests.incrementAndGet();
                    byte[] body = "ok".getBytes(StandardCharsets.US_ASCII);
                    String head = "HTTP/1.1 200 OK\r\nContent-Length: " + body.length
                            + "\r\nConnection: keep-alive\r\n\r\n";
                    out.write(head.getBytes(StandardCharsets.US_ASCII));
                    out.write(body);
                    out.flush();
                }
            } catch (IOException ignored) {
                // 客户端断开/连接被关闭，属预期
            } finally {
                sockets.remove(socket);
            }
        }

        /** 读取到空行分隔的请求头结束；EOF 返回 false。 */
        private static boolean readRequest(InputStream in) throws IOException {
            final byte[] end = {'\r', '\n', '\r', '\n'};
            int matched = 0;
            while (true) {
                int b = in.read();
                if (b == -1) {
                    return false;
                }
                if ((byte) b == end[matched]) {
                    matched++;
                } else {
                    matched = (b == '\r') ? 1 : 0;
                }
                if (matched == 4) {
                    return true;
                }
            }
        }

        @Override
        public void close() {
            running = false;
            try {
                serverSocket.close();
            } catch (IOException ignored) {
                // ignore
            }
            for (Socket socket : sockets) {
                try {
                    socket.close();
                } catch (IOException ignored) {
                    // ignore
                }
            }
        }
    }

    private static RestClient pooledClient() {
        return RestClient.builder()
                .requestFactory(HttpClientPool.requestFactory(5_000))
                .build();
    }

    private static void get(RestClient client, int port) {
        client.get().uri("http://127.0.0.1:" + port + "/").retrieve().toBodilessEntity();
    }

    @Test
    void sequentialRequestsReuseSingleConnection() throws Exception {
        try (CountingServer server = new CountingServer()) {
            RestClient client = pooledClient();
            for (int i = 0; i < 6; i++) {
                get(client, server.port());
            }
            assertEquals(6, server.requests());
            assertEquals(1, server.accepted(), "连续 6 次请求应复用同一连接");
        }
    }

    @Test
    void sharedAcrossClientFactoriesReuseOneConnection() throws Exception {
        try (CountingServer server = new CountingServer()) {
            // 三个独立 RestClient（各自 JdkClientHttpRequestFactory）共享底层连接池
            RestClient brand = pooledClient();
            RestClient ziniao = pooledClient();
            RestClient llm = pooledClient();
            get(brand, server.port());
            get(ziniao, server.port());
            get(llm, server.port());
            assertEquals(3, server.requests());
            assertEquals(1, server.accepted(), "不同客户端工厂共用同一连接池，应复用同一连接");
        }
    }

    @Test
    void reuseSurvivesShortIdleGap() throws Exception {
        try (CountingServer server = new CountingServer()) {
            RestClient client = pooledClient();
            for (int i = 0; i < 3; i++) {
                get(client, server.port());
            }
            Thread.sleep(80);
            for (int i = 0; i < 3; i++) {
                get(client, server.port());
            }
            assertEquals(6, server.requests());
            assertEquals(1, server.accepted(), "短暂空闲后连接应被 keep-alive 复用");
        }
    }

    @Test
    void manySequentialRequestsDoNotBurstConnections() throws Exception {
        try (CountingServer server = new CountingServer()) {
            RestClient client = pooledClient();
            for (int i = 0; i < 12; i++) {
                get(client, server.port());
            }
            assertEquals(12, server.requests());
            assertEquals(1, server.accepted(), "顺序突发不应新建连接");
        }
    }

    @Test
    void poolIsSingletonAcrossAccesses() {
        // 共享池单例；不同 RestClient 各自持工厂，但底层 HttpClient 唯一（Task77 已测实例级）
        assertSame(HttpClientPool.sharedHttpClient(), HttpClientPool.sharedHttpClient(),
                "共享池单例");
    }

    @Test
    void distinctHostsGetDistinctConnections() throws Exception {
        try (CountingServer a = new CountingServer(); CountingServer b = new CountingServer()) {
            RestClient client = pooledClient();
            get(client, a.port());
            get(client, b.port());
            get(client, a.port());
            assertEquals(1, a.accepted(), "host A 两次请求复用");
            assertEquals(1, b.accepted(), "host B 独立建连");
        }
    }

    @Test
    void nonPooledHttpClientDoesNotReusePoolConnection() throws Exception {
        try (CountingServer server = new CountingServer()) {
            RestClient pool = pooledClient();
            get(pool, server.port());
            assertEquals(1, server.accepted());

            // 独立 HttpClient（不参与共享池）访问同一 host → 新建连接，不复用池连接
            HttpClient standalone = HttpClient.newBuilder()
                    .connectTimeout(Duration.ofSeconds(5))
                    .build();
            org.springframework.http.client.JdkClientHttpRequestFactory factory =
                    new org.springframework.http.client.JdkClientHttpRequestFactory(standalone);
            RestClient standaloneRest = RestClient.builder().requestFactory(factory).build();
            get(standaloneRest, server.port());

            assertEquals(2, server.accepted(), "非共享池客户端应另建连接（不复用时暴露）");
        }
    }

    @Test
    void concurrentRequestsSucceedAndSettledConnectionsAreReused() throws Exception {
        try (CountingServer server = new CountingServer()) {
            RestClient client = pooledClient();
            int concurrency = 6;
            ExecutorService executor = Executors.newFixedThreadPool(concurrency);
            CountDownLatch start = new CountDownLatch(1);
            CountDownLatch done = new CountDownLatch(concurrency);
            try {
                IntStream.range(0, concurrency).forEach(i -> executor.submit(() -> {
                    try {
                        start.await();
                        get(client, server.port());
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    } finally {
                        done.countDown();
                    }
                }));
                start.countDown();
                assertTrue(done.await(10, TimeUnit.SECONDS), "并发请求应在时限内完成");
                assertEquals(concurrency, server.requests());
                int acceptedAfterConcurrent = server.accepted();
                assertTrue(acceptedAfterConcurrent >= 1 && acceptedAfterConcurrent <= concurrency,
                        "并发建连数应在合理范围，实际 " + acceptedAfterConcurrent);

                // 并发平息后追加顺序请求应复用已建立的空闲连接，不新增
                for (int i = 0; i < 3; i++) {
                    get(client, server.port());
                }
                assertEquals(concurrency + 3, server.requests());
                assertEquals(acceptedAfterConcurrent, server.accepted(),
                        "顺序追加应复用既有连接，不新增");
            } finally {
                executor.shutdownNow();
            }
        }
    }
}