1# Copyright 2020 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14import logging
15import signal
16
17from absl import app
18from absl import flags
19
20from bin.lib import common
21from framework import xds_flags
22from framework import xds_k8s_flags
23from framework.infrastructure import gcp
24from framework.infrastructure import k8s
25
26logger = logging.getLogger(__name__)
27# Flags
28_CMD = flags.DEFINE_enum('cmd',
29                         default='run',
30                         enum_values=['run', 'cleanup'],
31                         help='Command')
32_SECURE = flags.DEFINE_bool("secure",
33                            default=False,
34                            help="Run server in the secure mode")
35_REUSE_NAMESPACE = flags.DEFINE_bool("reuse_namespace",
36                                     default=True,
37                                     help="Use existing namespace if exists")
38_REUSE_SERVICE = flags.DEFINE_bool("reuse_service",
39                                   default=False,
40                                   help="Use existing service if exists")
41_FOLLOW = flags.DEFINE_bool("follow",
42                            default=False,
43                            help="Follow pod logs. "
44                            "Requires --collect_app_logs")
45_CLEANUP_NAMESPACE = flags.DEFINE_bool(
46    "cleanup_namespace",
47    default=False,
48    help="Delete namespace during resource cleanup")
49flags.adopt_module_key_flags(xds_flags)
50flags.adopt_module_key_flags(xds_k8s_flags)
51# Running outside of a test suite, so require explicit resource_suffix.
52flags.mark_flag_as_required("resource_suffix")
53
54
55def _make_sigint_handler(server_runner: common.KubernetesServerRunner):
56
57    def sigint_handler(sig, frame):
58        del sig, frame
59        print('Caught Ctrl+C. Shutting down the logs')
60        server_runner.stop_pod_dependencies(log_drain_sec=3)
61
62    return sigint_handler
63
64
65def main(argv):
66    if len(argv) > 1:
67        raise app.UsageError('Too many command-line arguments.')
68
69    # Must be called before KubernetesApiManager or GcpApiManager init.
70    xds_flags.set_socket_default_timeout_from_flag()
71
72    should_follow_logs = _FOLLOW.value and xds_flags.COLLECT_APP_LOGS.value
73    should_port_forward = (should_follow_logs and
74                           xds_k8s_flags.DEBUG_USE_PORT_FORWARDING.value)
75
76    # Setup.
77    gcp_api_manager = gcp.api.GcpApiManager()
78    k8s_api_manager = k8s.KubernetesApiManager(xds_k8s_flags.KUBE_CONTEXT.value)
79    server_namespace = common.make_server_namespace(k8s_api_manager)
80    server_runner = common.make_server_runner(
81        server_namespace,
82        gcp_api_manager,
83        reuse_namespace=_REUSE_NAMESPACE.value,
84        reuse_service=_REUSE_SERVICE.value,
85        secure=_SECURE.value,
86        port_forwarding=should_port_forward)
87
88    if _CMD.value == 'run':
89        logger.info('Run server, secure_mode=%s', _SECURE.value)
90        server_runner.run(
91            test_port=xds_flags.SERVER_PORT.value,
92            maintenance_port=xds_flags.SERVER_MAINTENANCE_PORT.value,
93            secure_mode=_SECURE.value,
94            log_to_stdout=_FOLLOW.value)
95        if should_follow_logs:
96            print('Following pod logs. Press Ctrl+C top stop')
97            signal.signal(signal.SIGINT, _make_sigint_handler(server_runner))
98            signal.pause()
99
100    elif _CMD.value == 'cleanup':
101        logger.info('Cleanup server')
102        server_runner.cleanup(force=True,
103                              force_namespace=_CLEANUP_NAMESPACE.value)
104
105
106if __name__ == '__main__':
107    app.run(main)
108