xref: /aosp_15_r20/external/grpc-grpc/examples/python/xds/server.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2)
1*cc02d7e2SAndroid Build Coastguard Worker# Copyright 2020 The gRPC authors.
2*cc02d7e2SAndroid Build Coastguard Worker#
3*cc02d7e2SAndroid Build Coastguard Worker# Licensed under the Apache License, Version 2.0 (the "License");
4*cc02d7e2SAndroid Build Coastguard Worker# you may not use this file except in compliance with the License.
5*cc02d7e2SAndroid Build Coastguard Worker# You may obtain a copy of the License at
6*cc02d7e2SAndroid Build Coastguard Worker#
7*cc02d7e2SAndroid Build Coastguard Worker#     http://www.apache.org/licenses/LICENSE-2.0
8*cc02d7e2SAndroid Build Coastguard Worker#
9*cc02d7e2SAndroid Build Coastguard Worker# Unless required by applicable law or agreed to in writing, software
10*cc02d7e2SAndroid Build Coastguard Worker# distributed under the License is distributed on an "AS IS" BASIS,
11*cc02d7e2SAndroid Build Coastguard Worker# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12*cc02d7e2SAndroid Build Coastguard Worker# See the License for the specific language governing permissions and
13*cc02d7e2SAndroid Build Coastguard Worker# limitations under the License.
14*cc02d7e2SAndroid Build Coastguard Worker"""The Python implementation of the GRPC helloworld.Greeter server."""
15*cc02d7e2SAndroid Build Coastguard Worker
16*cc02d7e2SAndroid Build Coastguard Workerimport argparse
17*cc02d7e2SAndroid Build Coastguard Workerfrom concurrent import futures
18*cc02d7e2SAndroid Build Coastguard Workerimport logging
19*cc02d7e2SAndroid Build Coastguard Workerimport socket
20*cc02d7e2SAndroid Build Coastguard Worker
21*cc02d7e2SAndroid Build Coastguard Workerimport grpc
22*cc02d7e2SAndroid Build Coastguard Workerfrom grpc_health.v1 import health
23*cc02d7e2SAndroid Build Coastguard Workerfrom grpc_health.v1 import health_pb2
24*cc02d7e2SAndroid Build Coastguard Workerfrom grpc_health.v1 import health_pb2_grpc
25*cc02d7e2SAndroid Build Coastguard Workerfrom grpc_reflection.v1alpha import reflection
26*cc02d7e2SAndroid Build Coastguard Workerimport helloworld_pb2
27*cc02d7e2SAndroid Build Coastguard Workerimport helloworld_pb2_grpc
28*cc02d7e2SAndroid Build Coastguard Worker
29*cc02d7e2SAndroid Build Coastguard Worker_DESCRIPTION = "A general purpose phony server."
30*cc02d7e2SAndroid Build Coastguard Worker
31*cc02d7e2SAndroid Build Coastguard Worker_LISTEN_HOST = "0.0.0.0"
32*cc02d7e2SAndroid Build Coastguard Worker
33*cc02d7e2SAndroid Build Coastguard Worker_THREAD_POOL_SIZE = 256
34*cc02d7e2SAndroid Build Coastguard Worker
35*cc02d7e2SAndroid Build Coastguard Workerlogger = logging.getLogger()
36*cc02d7e2SAndroid Build Coastguard Workerconsole_handler = logging.StreamHandler()
37*cc02d7e2SAndroid Build Coastguard Workerformatter = logging.Formatter(fmt="%(asctime)s: %(levelname)-8s %(message)s")
38*cc02d7e2SAndroid Build Coastguard Workerconsole_handler.setFormatter(formatter)
39*cc02d7e2SAndroid Build Coastguard Workerlogger.addHandler(console_handler)
40*cc02d7e2SAndroid Build Coastguard Worker
41*cc02d7e2SAndroid Build Coastguard Worker
42*cc02d7e2SAndroid Build Coastguard Workerclass Greeter(helloworld_pb2_grpc.GreeterServicer):
43*cc02d7e2SAndroid Build Coastguard Worker    def __init__(self, hostname: str):
44*cc02d7e2SAndroid Build Coastguard Worker        self._hostname = hostname if hostname else socket.gethostname()
45*cc02d7e2SAndroid Build Coastguard Worker
46*cc02d7e2SAndroid Build Coastguard Worker    def SayHello(
47*cc02d7e2SAndroid Build Coastguard Worker        self,
48*cc02d7e2SAndroid Build Coastguard Worker        request: helloworld_pb2.HelloRequest,
49*cc02d7e2SAndroid Build Coastguard Worker        context: grpc.ServicerContext,
50*cc02d7e2SAndroid Build Coastguard Worker    ) -> helloworld_pb2.HelloReply:
51*cc02d7e2SAndroid Build Coastguard Worker        return helloworld_pb2.HelloReply(
52*cc02d7e2SAndroid Build Coastguard Worker            message=f"Hello {request.name} from {self._hostname}!"
53*cc02d7e2SAndroid Build Coastguard Worker        )
54*cc02d7e2SAndroid Build Coastguard Worker
55*cc02d7e2SAndroid Build Coastguard Worker
56*cc02d7e2SAndroid Build Coastguard Workerdef _configure_maintenance_server(
57*cc02d7e2SAndroid Build Coastguard Worker    server: grpc.Server, maintenance_port: int
58*cc02d7e2SAndroid Build Coastguard Worker) -> None:
59*cc02d7e2SAndroid Build Coastguard Worker    listen_address = f"{_LISTEN_HOST}:{maintenance_port}"
60*cc02d7e2SAndroid Build Coastguard Worker    server.add_insecure_port(listen_address)
61*cc02d7e2SAndroid Build Coastguard Worker
62*cc02d7e2SAndroid Build Coastguard Worker    # Create a health check servicer. We use the non-blocking implementation
63*cc02d7e2SAndroid Build Coastguard Worker    # to avoid thread starvation.
64*cc02d7e2SAndroid Build Coastguard Worker    health_servicer = health.HealthServicer(
65*cc02d7e2SAndroid Build Coastguard Worker        experimental_non_blocking=True,
66*cc02d7e2SAndroid Build Coastguard Worker        experimental_thread_pool=futures.ThreadPoolExecutor(
67*cc02d7e2SAndroid Build Coastguard Worker            max_workers=_THREAD_POOL_SIZE
68*cc02d7e2SAndroid Build Coastguard Worker        ),
69*cc02d7e2SAndroid Build Coastguard Worker    )
70*cc02d7e2SAndroid Build Coastguard Worker
71*cc02d7e2SAndroid Build Coastguard Worker    # Create a tuple of all of the services we want to export via reflection.
72*cc02d7e2SAndroid Build Coastguard Worker    services = tuple(
73*cc02d7e2SAndroid Build Coastguard Worker        service.full_name
74*cc02d7e2SAndroid Build Coastguard Worker        for service in helloworld_pb2.DESCRIPTOR.services_by_name.values()
75*cc02d7e2SAndroid Build Coastguard Worker    ) + (reflection.SERVICE_NAME, health.SERVICE_NAME)
76*cc02d7e2SAndroid Build Coastguard Worker
77*cc02d7e2SAndroid Build Coastguard Worker    # Mark all services as healthy.
78*cc02d7e2SAndroid Build Coastguard Worker    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
79*cc02d7e2SAndroid Build Coastguard Worker    for service in services:
80*cc02d7e2SAndroid Build Coastguard Worker        health_servicer.set(service, health_pb2.HealthCheckResponse.SERVING)
81*cc02d7e2SAndroid Build Coastguard Worker    reflection.enable_server_reflection(services, server)
82*cc02d7e2SAndroid Build Coastguard Worker
83*cc02d7e2SAndroid Build Coastguard Worker
84*cc02d7e2SAndroid Build Coastguard Workerdef _configure_greeter_server(
85*cc02d7e2SAndroid Build Coastguard Worker    server: grpc.Server, port: int, secure_mode: bool, hostname
86*cc02d7e2SAndroid Build Coastguard Worker) -> None:
87*cc02d7e2SAndroid Build Coastguard Worker    # Add the application servicer to the server.
88*cc02d7e2SAndroid Build Coastguard Worker    helloworld_pb2_grpc.add_GreeterServicer_to_server(Greeter(hostname), server)
89*cc02d7e2SAndroid Build Coastguard Worker    listen_address = f"{_LISTEN_HOST}:{port}"
90*cc02d7e2SAndroid Build Coastguard Worker    if not secure_mode:
91*cc02d7e2SAndroid Build Coastguard Worker        server.add_insecure_port(listen_address)
92*cc02d7e2SAndroid Build Coastguard Worker    else:
93*cc02d7e2SAndroid Build Coastguard Worker        # Use xDS credentials.
94*cc02d7e2SAndroid Build Coastguard Worker        logger.info("Running with xDS Server credentials")
95*cc02d7e2SAndroid Build Coastguard Worker
96*cc02d7e2SAndroid Build Coastguard Worker        # Fall back to insecure credentials.
97*cc02d7e2SAndroid Build Coastguard Worker        server_fallback_creds = grpc.insecure_server_credentials()
98*cc02d7e2SAndroid Build Coastguard Worker        server_creds = grpc.xds_server_credentials(server_fallback_creds)
99*cc02d7e2SAndroid Build Coastguard Worker        server.add_secure_port(listen_address, server_creds)
100*cc02d7e2SAndroid Build Coastguard Worker
101*cc02d7e2SAndroid Build Coastguard Worker
102*cc02d7e2SAndroid Build Coastguard Workerdef serve(
103*cc02d7e2SAndroid Build Coastguard Worker    port: int, hostname: str, maintenance_port: int, secure_mode: bool
104*cc02d7e2SAndroid Build Coastguard Worker) -> None:
105*cc02d7e2SAndroid Build Coastguard Worker    if port == maintenance_port:
106*cc02d7e2SAndroid Build Coastguard Worker        # If maintenance port and port are the same, start a single server.
107*cc02d7e2SAndroid Build Coastguard Worker        server = grpc.server(
108*cc02d7e2SAndroid Build Coastguard Worker            futures.ThreadPoolExecutor(max_workers=_THREAD_POOL_SIZE)
109*cc02d7e2SAndroid Build Coastguard Worker        )
110*cc02d7e2SAndroid Build Coastguard Worker        _configure_greeter_server(server, port, secure_mode, hostname)
111*cc02d7e2SAndroid Build Coastguard Worker        _configure_maintenance_server(server, maintenance_port)
112*cc02d7e2SAndroid Build Coastguard Worker        server.start()
113*cc02d7e2SAndroid Build Coastguard Worker        logger.info("Greeter server listening on port %d", port)
114*cc02d7e2SAndroid Build Coastguard Worker        logger.info("Maintenance server listening on port %d", maintenance_port)
115*cc02d7e2SAndroid Build Coastguard Worker        server.wait_for_termination()
116*cc02d7e2SAndroid Build Coastguard Worker    else:
117*cc02d7e2SAndroid Build Coastguard Worker        # Otherwise, start two different servers.
118*cc02d7e2SAndroid Build Coastguard Worker        greeter_server = grpc.server(
119*cc02d7e2SAndroid Build Coastguard Worker            futures.ThreadPoolExecutor(max_workers=_THREAD_POOL_SIZE),
120*cc02d7e2SAndroid Build Coastguard Worker            xds=secure_mode,
121*cc02d7e2SAndroid Build Coastguard Worker        )
122*cc02d7e2SAndroid Build Coastguard Worker        _configure_greeter_server(greeter_server, port, secure_mode, hostname)
123*cc02d7e2SAndroid Build Coastguard Worker        greeter_server.start()
124*cc02d7e2SAndroid Build Coastguard Worker        logger.info("Greeter server listening on port %d", port)
125*cc02d7e2SAndroid Build Coastguard Worker        maintenance_server = grpc.server(
126*cc02d7e2SAndroid Build Coastguard Worker            futures.ThreadPoolExecutor(max_workers=_THREAD_POOL_SIZE)
127*cc02d7e2SAndroid Build Coastguard Worker        )
128*cc02d7e2SAndroid Build Coastguard Worker        _configure_maintenance_server(maintenance_server, maintenance_port)
129*cc02d7e2SAndroid Build Coastguard Worker        maintenance_server.start()
130*cc02d7e2SAndroid Build Coastguard Worker        logger.info("Maintenance server listening on port %d", maintenance_port)
131*cc02d7e2SAndroid Build Coastguard Worker        greeter_server.wait_for_termination()
132*cc02d7e2SAndroid Build Coastguard Worker        maintenance_server.wait_for_termination()
133*cc02d7e2SAndroid Build Coastguard Worker
134*cc02d7e2SAndroid Build Coastguard Worker
135*cc02d7e2SAndroid Build Coastguard Workerif __name__ == "__main__":
136*cc02d7e2SAndroid Build Coastguard Worker    parser = argparse.ArgumentParser(description=_DESCRIPTION)
137*cc02d7e2SAndroid Build Coastguard Worker    parser.add_argument(
138*cc02d7e2SAndroid Build Coastguard Worker        "port",
139*cc02d7e2SAndroid Build Coastguard Worker        default=50051,
140*cc02d7e2SAndroid Build Coastguard Worker        type=int,
141*cc02d7e2SAndroid Build Coastguard Worker        nargs="?",
142*cc02d7e2SAndroid Build Coastguard Worker        help="The port on which to listen.",
143*cc02d7e2SAndroid Build Coastguard Worker    )
144*cc02d7e2SAndroid Build Coastguard Worker    parser.add_argument(
145*cc02d7e2SAndroid Build Coastguard Worker        "hostname",
146*cc02d7e2SAndroid Build Coastguard Worker        type=str,
147*cc02d7e2SAndroid Build Coastguard Worker        default=None,
148*cc02d7e2SAndroid Build Coastguard Worker        nargs="?",
149*cc02d7e2SAndroid Build Coastguard Worker        help="The name clients will see in responses.",
150*cc02d7e2SAndroid Build Coastguard Worker    )
151*cc02d7e2SAndroid Build Coastguard Worker    parser.add_argument(
152*cc02d7e2SAndroid Build Coastguard Worker        "--xds-creds",
153*cc02d7e2SAndroid Build Coastguard Worker        action="store_true",
154*cc02d7e2SAndroid Build Coastguard Worker        help="If specified, uses xDS credentials to connect to the server.",
155*cc02d7e2SAndroid Build Coastguard Worker    )
156*cc02d7e2SAndroid Build Coastguard Worker    args = parser.parse_args()
157*cc02d7e2SAndroid Build Coastguard Worker    logging.basicConfig()
158*cc02d7e2SAndroid Build Coastguard Worker    logger.setLevel(logging.INFO)
159*cc02d7e2SAndroid Build Coastguard Worker    serve(args.port, args.hostname, args.port + 1, args.xds_creds)
160