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 client in the secure mode")
35_QPS = flags.DEFINE_integer('qps', default=25, help='Queries per second')
36_PRINT_RESPONSE = flags.DEFINE_bool("print_response",
37                                    default=False,
38                                    help="Client prints responses")
39_FOLLOW = flags.DEFINE_bool(
40    "follow",
41    default=False,
42    help=
43    "Follow pod logs. Requires --collect_app_logs or --debug_use_port_forwarding"
44)
45_CONFIG_MESH = flags.DEFINE_bool(
46    "config_mesh",
47    default=None,
48    help="Optional. Supplied to bootstrap generator to indicate AppNet mesh.")
49_REUSE_NAMESPACE = flags.DEFINE_bool("reuse_namespace",
50                                     default=True,
51                                     help="Use existing namespace if exists")
52_CLEANUP_NAMESPACE = flags.DEFINE_bool(
53    "cleanup_namespace",
54    default=False,
55    help="Delete namespace during resource cleanup")
56flags.adopt_module_key_flags(xds_flags)
57flags.adopt_module_key_flags(xds_k8s_flags)
58# Running outside of a test suite, so require explicit resource_suffix.
59flags.mark_flag_as_required(xds_flags.RESOURCE_SUFFIX.name)
60
61
62@flags.multi_flags_validator((xds_flags.SERVER_XDS_PORT.name, _CMD.name),
63                             message="Run outside of a test suite, must provide"
64                             " the exact port value (must be greater than 0).")
65def _check_server_xds_port_flag(flags_dict):
66    if flags_dict[_CMD.name] == "cleanup":
67        return True
68    return flags_dict[xds_flags.SERVER_XDS_PORT.name] > 0
69
70
71def _make_sigint_handler(client_runner: common.KubernetesClientRunner):
72
73    def sigint_handler(sig, frame):
74        del sig, frame
75        print('Caught Ctrl+C. Shutting down the logs')
76        client_runner.stop_pod_dependencies(log_drain_sec=3)
77
78    return sigint_handler
79
80
81def main(argv):
82    if len(argv) > 1:
83        raise app.UsageError('Too many command-line arguments.')
84
85    # Must be called before KubernetesApiManager or GcpApiManager init.
86    xds_flags.set_socket_default_timeout_from_flag()
87
88    # Log following and port forwarding.
89    should_follow_logs = _FOLLOW.value and xds_flags.COLLECT_APP_LOGS.value
90    should_port_forward = (should_follow_logs and
91                           xds_k8s_flags.DEBUG_USE_PORT_FORWARDING.value)
92
93    # Setup.
94    gcp_api_manager = gcp.api.GcpApiManager()
95    k8s_api_manager = k8s.KubernetesApiManager(xds_k8s_flags.KUBE_CONTEXT.value)
96    client_namespace = common.make_client_namespace(k8s_api_manager)
97    client_runner = common.make_client_runner(
98        client_namespace,
99        gcp_api_manager,
100        reuse_namespace=_REUSE_NAMESPACE.value,
101        secure=_SECURE.value,
102        port_forwarding=should_port_forward)
103
104    # Server target
105    server_xds_host = xds_flags.SERVER_XDS_HOST.value
106    server_xds_port = xds_flags.SERVER_XDS_PORT.value
107
108    if _CMD.value == 'run':
109        logger.info('Run client, secure_mode=%s', _SECURE.value)
110        client_runner.run(
111            server_target=f'xds:///{server_xds_host}:{server_xds_port}',
112            qps=_QPS.value,
113            print_response=_PRINT_RESPONSE.value,
114            secure_mode=_SECURE.value,
115            config_mesh=_CONFIG_MESH.value,
116            log_to_stdout=_FOLLOW.value)
117        if should_follow_logs:
118            print('Following pod logs. Press Ctrl+C top stop')
119            signal.signal(signal.SIGINT, _make_sigint_handler(client_runner))
120            signal.pause()
121
122    elif _CMD.value == 'cleanup':
123        logger.info('Cleanup client')
124        client_runner.cleanup(force=True,
125                              force_namespace=_CLEANUP_NAMESPACE.value)
126
127
128if __name__ == '__main__':
129    app.run(main)
130