xref: /aosp_15_r20/external/federated-compute/fcp/client/http/testing/http_test_server.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2022 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "fcp/client/http/testing/http_test_server.h"
18 
19 #include <functional>
20 #include <memory>
21 #include <string>
22 #include <utility>
23 
24 #include "absl/status/statusor.h"
25 #include "fcp/base/scheduler.h"
26 
27 namespace fcp {
28 namespace client {
29 namespace http {
30 namespace {
31 using ::tensorflow::serving::net_http::EventExecutor;
32 using ::tensorflow::serving::net_http::HTTPServerInterface;
33 using ::tensorflow::serving::net_http::RequestHandlerOptions;
34 using ::tensorflow::serving::net_http::ServerOptions;
35 using ::tensorflow::serving::net_http::ServerRequestInterface;
36 
37 // Echoes the request back to the client
EchoHandler(ServerRequestInterface * req)38 void EchoHandler(ServerRequestInterface* req) {
39   std::string response;
40 
41   absl::StrAppend(&response, "HTTP Method: ", req->http_method(), "\n");
42   absl::StrAppend(&response, "Request Uri: ", req->uri_path(), "\n");
43 
44   absl::StrAppend(&response, "Request Headers:\n");
45   for (absl::string_view header : req->request_headers()) {
46     absl::StrAppend(&response, header, ": ", req->GetRequestHeader(header),
47                     "\n");
48   }
49 
50   absl::StrAppend(&response, "Request Body:\n");
51   // Read the request body
52   int64_t num_bytes;
53   while (true) {
54     auto request_chunk = req->ReadRequestBytes(&num_bytes);
55     if (request_chunk == nullptr) {
56       break;
57     }
58     absl::StrAppend(&response,
59                     absl::string_view(request_chunk.get(), num_bytes));
60   }
61 
62   req->WriteResponseString(response);
63 
64   SetContentTypeHTML(req);
65   req->Reply();
66 }
67 
68 // Non-blocking event executor needed for an event-driven web server
69 class ThreadPoolEventExecutor final : public EventExecutor {
70  public:
ThreadPoolEventExecutor(int num_threads)71   explicit ThreadPoolEventExecutor(int num_threads)
72       : thread_pool_scheduler_(CreateThreadPoolScheduler(num_threads)) {}
~ThreadPoolEventExecutor()73   ~ThreadPoolEventExecutor() override {
74     thread_pool_scheduler_->WaitUntilIdle();
75   }
76 
Schedule(std::function<void ()> fn)77   void Schedule(std::function<void()> fn) override {
78     thread_pool_scheduler_->Schedule(fn);
79   }
80 
81  private:
82   std::unique_ptr<Scheduler> thread_pool_scheduler_;
83 };
84 }  // namespace
85 
CreateHttpTestServer(const std::string & uri,int port,int num_threads)86 absl::StatusOr<std::unique_ptr<HTTPServerInterface>> CreateHttpTestServer(
87     const std::string& uri, int port, int num_threads) {
88   auto options = std::make_unique<ServerOptions>();
89   options->AddPort(port);
90   options->SetExecutor(std::make_unique<ThreadPoolEventExecutor>(num_threads));
91 
92   std::unique_ptr<HTTPServerInterface> http_server =
93       CreateEvHTTPServer(std::move(options));
94   if (http_server == nullptr) {
95     return absl::InternalError("Failed to create EvHTTPServer");
96   }
97 
98   RequestHandlerOptions handler_options;
99   http_server->RegisterRequestHandler(uri, EchoHandler, handler_options);
100   return http_server;
101 }
102 
103 }  // namespace http
104 }  // namespace client
105 }  // namespace fcp
106