基於Spring的流量拷貝框架實現

1 背景

目前我們在開發一個大資料線上查詢系統服務，該服務下面會支援多種資料庫引擎的查詢，比如Impala、Kylin和Druid等，並根據查詢請求進行自動路由，選擇最優的資料庫引擎。因此路由演算法是決定整體查詢效能好壞的關鍵。原本線上採用的是按照經驗的資料庫預設排序規則，現在我們開發了一套基於神經網路的預測模型，想對這兩種方式的效能指標進行對比。

由於我們必須要保證線上服務的穩定性，不能將一部分流量切到測試環境，因此這是一個非典型的A/B Test的問題。

我們最終考慮的方案，是仍然將全量請求傳送到線上環境，同時將一定比例的請求拷貝到測試環境，再將這部分資料進行One on one的效能對比，以判斷新模型的效能優化效果。

由於公司部署環境的限制，是不能採用類似tcpcopy或者gor之類的工具進行流量拷貝的。在考慮自己實現時，因為我們資料查詢系統是基於Spring Cloud+Spring Boot的微服務框架開發的， 那麼如何才能快速、方便、並對現有程式碼減少耦合的方式來實現本功能，是本文要討論的內容。

2 實現思路

該功能的實現思路並不複雜。在Controller接收到請求後，新建立一個執行緒，向指定的測試環境地址傳送一個引數完全相同的請求即可。

2.1 簡單實現

根據上述的描述，可以簡單的實現以下程式碼，將20%的請求傳送到url-test：

@RequestMapping(value = "/test/post", method = RequestMethod.POST)
public String helloTest(@RequestParam String value, @RequestBody PostBody postBody) {
  if (new Random().nextFloat() < 0.2F) {
    new Thread(() -> restTemplate.postForObject("http://url-test?value=" + value, postBody, String.class)).start();
  }
  return "hello test, " + postBody;
}

但是該程式碼的可複用性極差，不僅需要侵入@Controller每個方法的程式碼，並且無法控制該功能是否關閉，同時由於無法獲取@RequestParam的引數名稱，還需要以字串形式寫在url中，很難維護。

2.2 優化實現

由於這裡要實現的功能還是比較明確的，所以我們考慮能不能將這部分程式碼抽離出來，然後以一種統一的方式進行配置，這裡就很自然考慮到Spring的AOP原理，那麼需要使用到自定義註解。

3 註解設計

採用統一的註解，標註在需要執行流量拷貝的方法上，是一種優雅的不侵入現有程式碼的處理方式。註解需要定義的變數包含以下兩個：

• 流量拷貝後傳送的目的地址URL
• 流量拷貝的比例，預設為10%

該註解定義的程式碼如下：

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestCopy {
  String url();
  float ratio() default 0.1F;
}

新增該註解後的Controller方法如下所示：

  @RequestMapping(value = "/post", method = RequestMethod.POST)
  @RequestCopy(url = "http://localhost:8080/test", ratio = 1F)
  public String hello(@RequestBody PostBody postBody) {
    System.out.println("hello, " + postBody);
    return "hello, " + postBody;
  }

4 註解處理

在定義好註解之後，選擇合適的時機獲取被註解的方法，並處理流量拷貝邏輯就是下一個問題。本專案是基於Spring Boot開發的，那Spring MVC的攔截器（interceptor）就是一個不錯的選擇。

攔截器：Spring MVC中的攔截器（Interceptor）類似於Servlet中的過濾器（Filter），它主要用於攔截使用者請求並作相應的處理。例如通過攔截器可以進行許可權驗證、記錄請求資訊的日誌、判斷使用者是否登入等。

也就是說，在攔截到被@RequestCopy註解修飾的方法後，通過獲取註解中的URL和ratio引數，將請求傳送到該地址，同時繼續將請求傳送到@Controller的方法。

這裡不得不說，攔截器提供的輸入引數非常有用，分別進行一下說明：

• handler：可通過強轉轉為Method類，從而獲取method上的所有註解，並判斷是否包含@RequestCopy註解，僅對包含@RequestCopy註解的方法進行後續處理。
• request：可通過getQueryString()方法將URL中root地址後的所有引數都獲取到，解決了引數名獲取的問題。

在獲取到之後，進行流量傳送的程式碼就是這裡的業務邏輯程式碼，這裡我們的場景是支援Get請求和Post json請求，那也就在程式碼中針對這兩種情況進行處理。

@Override
public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
    throws Exception {
  if (!(handler instanceof HandlerMethod)) {
    return true;
  }

  HandlerMethod handlerMethod = (HandlerMethod) handler;
  Method method = handlerMethod.getMethod();
  RequestCopy requestCopy = method.getAnnotation(RequestCopy.class);
  if (requestCopy == null) {
    return true;
  }

  String url = requestCopy.url() + request.getRequestURI();
  float ratio = requestCopy.ratio();
  if (new Random().nextFloat() <= ratio) {
    switch (request.getMethod()) {
      case "GET": {
        new Thread(() -> {
          String fullUrl = url + (request.getQueryString() == null ? "" : "?" + request.getQueryString());
          String result = restTemplate.getForObject(fullUrl, String.class);
          logger.info("Send copied GET request to url: {}, and receive response: {}", fullUrl, result);
        }).start();
        break;
      }
      case "POST": {
        switch (request.getHeader("Content-Type")) {
          case "application/json": {
            RequestBodyWrapper requestWrapper = new RequestBodyWrapper(request);
            JSONObject jsonObject = JSON.parseObject(requestWrapper.getBody());
            if (jsonObject != null) {
              new Thread(() -> {
                String fullUrl = url + (request.getQueryString() == null ? "" : "?" + request.getQueryString());
                String result = restTemplate.postForObject(fullUrl, jsonObject, String.class);
                logger.info("Send copied POST request to url: {}, body: {}, and receive response: {}", fullUrl,
                    jsonObject, result);
              }).start();
            }
            break;
          }
        }
        break;
      }
    }
  }
  return true;
}

5 Post方法的特殊處理

上述程式碼中，對Post方法的請求進行了一次包裝。這主要是因為，Post方法中的RequestBody是以資料流的形式傳輸的，在interceptor中獲取到RequestBody併傳送請求後，在真正進入@Controller的方法中時，就沒有RequestBody的值了。

因此這裡通過Servlet的Filter方法來解決此問題。在服務接收到請求後，先將RequestBody的內容儲存在Wrapper類的變數中，並重寫對應的getInputStream等方法，以實現RequestBody內容的重複讀取。

Filter也稱之為過濾器，它是Servlet技術中最實用的技術，WEB開發人員通過Filter技術，對web伺服器管理的所有web資源：例如Jsp, Servlet, 靜態圖片檔案或靜態 html 檔案等進行攔截，從而實現一些特殊的功能。例如實現URL級別的許可權訪問控制、過濾敏感詞彙、壓縮響應資訊等一些高階功能。

Servlet的Filter方法的執行時機，是一定會在Spring MVC的interceptor之前的，因此也可以滿足在接收到請求的第一時間對該請求進行包裝的需求。

public class RequestBodyWrapper extends HttpServletRequestWrapper {

  private final String body;

  public RequestBodyWrapper(HttpServletRequest request) throws IOException {
    super(request);
    StringBuilder stringBuilder = new StringBuilder();
    BufferedReader bufferedReader = null;
    try {
      InputStream inputStream = request.getInputStream();
      if (inputStream != null) {
        bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
        char[] charBuffer = new char[128];
        int bytesRead;
        while ((bytesRead = bufferedReader.read(charBuffer)) > 0) {
          stringBuilder.append(charBuffer, 0, bytesRead);
        }
      }
    } finally {
      if (bufferedReader != null) {
        bufferedReader.close();
      }
    }
    body = stringBuilder.toString();
  }

  @Override
  public ServletInputStream getInputStream() {
    final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(body.getBytes());
    return new ServletInputStream() {
      @Override
      public boolean isFinished() {
        return false;
      }

      @Override
      public boolean isReady() {
        return false;
      }

      @Override
      public void setReadListener(ReadListener readListener) {

      }

      public int read() {
        return byteArrayInputStream.read();
      }
    };
  }

  @Override
  public BufferedReader getReader() {
    return new BufferedReader(new InputStreamReader(this.getInputStream()));
  }

  public String getBody() {
    return this.body;
  }
}

6 配置開關

上述的程式碼開發完成後，還需要將interceptor和filter進行註冊，註冊的方式就是進行相對配置介面的重寫，具體可參考最後的程式碼示例。

這裡通過使用@ConditionOnProperty註解，可以在Spring Boot的application.yml配置檔案中，通過進行配置項值的修改，來決定是否開啟該功能。

7 使用說明

在該專案完成後，我們將之打包並上傳到Maven中央倉庫中。在使用時，需要先載入兩個重寫的配置類，然後在@Controller的方法上增加@RequestCopy註解，最後在application.yml中增加request-copy: true開啟該功能，即完成了對該方法的流量拷貝的功能。

8 總結

在實際工作中遇到可以抽象化的功能時，要儘量實現功能的模組化。這裡採用了Java的自定義註解、Spring MVC的攔截器和Servlet的過濾器三個特性，在實現此功能的同時，也加深了我們會這三個功能的理解，希望和大家一起學習和交流。

9 專案地址

• Github: https://github.com/meazza/spring-boot-request-copy
• Maven: https://mvnrepository.com/artifact/com.github.meazza/spring-boot-request-copy