ceph/src/jaegertracing/thrift/contrib/fb303/py/fb303_scripts/fb303_simple_mgmt.py

   1 #!/usr/bin/env python
   2
   3 #
   4 # Licensed to the Apache Software Foundation (ASF) under one
   5 # or more contributor license agreements. See the NOTICE file
   6 # distributed with this work for additional information
   7 # regarding copyright ownership. The ASF licenses this file
   8 # to you under the Apache License, Version 2.0 (the
   9 # "License"); you may not use this file except in compliance
  10 # with the License. You may obtain a copy of the License at
  11 #
  12 #   http://www.apache.org/licenses/LICENSE-2.0
  13 #
  14 # Unless required by applicable law or agreed to in writing,
  15 # software distributed under the License is distributed on an
  16 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  17 # KIND, either express or implied. See the License for the
  18 # specific language governing permissions and limitations
  19 # under the License.
  20 #
  21
  22 from __future__ import print_function
  23 import sys
  24 import os
  25 from optparse import OptionParser
  26
  27 from thrift.Thrift import *
  28
  29 from thrift.transport import TSocket
  30 from thrift.transport import TTransport
  31 from thrift.protocol import TBinaryProtocol
  32
  33 from fb303 import *
  34 from fb303.ttypes import *
  35
  36
  37 def service_ctrl(
  38         command,
  39         port,
  40         trans_factory=None,
  41         prot_factory=None):
  42     """
  43     service_ctrl is a generic function to execute standard fb303 functions
  44
  45     @param command: one of stop, start, reload, status, counters, name, alive
  46     @param port: service's port
  47     @param trans_factory: TTransportFactory to use for obtaining a TTransport. Default is
  48                           TBufferedTransportFactory
  49     @param prot_factory: TProtocolFactory to use for obtaining a TProtocol. Default is
  50                          TBinaryProtocolFactory
  51     """
  52
  53     if command in ["status"]:
  54         try:
  55             status = fb303_wrapper('status', port, trans_factory, prot_factory)
  56             status_details = fb303_wrapper('get_status_details', port, trans_factory, prot_factory)
  57
  58             msg = fb_status_string(status)
  59             if (len(status_details)):
  60                 msg += " - %s" % status_details
  61             print(msg)
  62             return 2 if status == fb_status.ALIVE else 3
  63         except:
  64             print("Failed to get status")
  65             return 3
  66
  67     # scalar commands
  68     if command in ["version", "alive", "name"]:
  69         try:
  70             result = fb303_wrapper(command, port, trans_factory, prot_factory)
  71             print(result)
  72             return 0
  73         except:
  74             print("failed to get ", command)
  75             return 3
  76
  77     # counters
  78     if command in ["counters"]:
  79         try:
  80             counters = fb303_wrapper('counters', port, trans_factory, prot_factory)
  81             for counter in counters:
  82                 print("%s: %d" % (counter.encode('utf-8'), counters[counter]))
  83             return 0
  84         except:
  85             print("failed to get counters")
  86             return 3
  87
  88     # Only root should be able to run the following commands
  89     if os.getuid() == 0:
  90         # async commands
  91         if command in ["stop", "reload"]:
  92             try:
  93                 fb303_wrapper(command, port, trans_factory, prot_factory)
  94                 return 0
  95             except:
  96                 print("failed to tell the service to ", command)
  97                 return 3
  98     else:
  99         if command in ["stop", "reload"]:
 100             print("root privileges are required to stop or reload the service.")
 101             return 4
 102
 103     print("The following commands are available:")
 104     for command in ["counters", "name", "version", "alive", "status"]:
 105         print("\t%s" % command)
 106     print("The following commands are available for users with root privileges:")
 107     for command in ["stop", "reload"]:
 108         print("\t%s" % command)
 109
 110     return 0
 111
 112
 113 def fb303_wrapper(command, port, trans_factory=None, prot_factory=None):
 114     sock = TSocket.TSocket('localhost', port)
 115
 116     # use input transport factory if provided
 117     if (trans_factory is None):
 118         trans = TTransport.TBufferedTransport(sock)
 119     else:
 120         trans = trans_factory.getTransport(sock)
 121
 122     # use input protocol factory if provided
 123     if (prot_factory is None):
 124         prot = TBinaryProtocol.TBinaryProtocol(trans)
 125     else:
 126         prot = prot_factory.getProtocol(trans)
 127
 128     # initialize client and open transport
 129     fb303_client = FacebookService.Client(prot, prot)
 130     trans.open()
 131
 132     if (command == 'reload'):
 133         fb303_client.reinitialize()
 134
 135     elif (command == 'stop'):
 136         fb303_client.shutdown()
 137
 138     elif (command == 'status'):
 139         return fb303_client.getStatus()
 140
 141     elif (command == 'version'):
 142         return fb303_client.getVersion()
 143
 144     elif (command == 'get_status_details'):
 145         return fb303_client.getStatusDetails()
 146
 147     elif (command == 'counters'):
 148         return fb303_client.getCounters()
 149
 150     elif (command == 'name'):
 151         return fb303_client.getName()
 152
 153     elif (command == 'alive'):
 154         return fb303_client.aliveSince()
 155
 156     trans.close()
 157
 158
 159 def fb_status_string(status_enum):
 160     if (status_enum == fb_status.DEAD):
 161         return "DEAD"
 162     if (status_enum == fb_status.STARTING):
 163         return "STARTING"
 164     if (status_enum == fb_status.ALIVE):
 165         return "ALIVE"
 166     if (status_enum == fb_status.STOPPING):
 167         return "STOPPING"
 168     if (status_enum == fb_status.STOPPED):
 169         return "STOPPED"
 170     if (status_enum == fb_status.WARNING):
 171         return "WARNING"
 172
 173
 174 def main():
 175
 176     # parse command line options
 177     parser = OptionParser()
 178     commands = ["stop", "counters", "status", "reload", "version", "name", "alive"]
 179
 180     parser.add_option("-c", "--command", dest="command", help="execute this API",
 181                       choices=commands, default="status")
 182     parser.add_option("-p", "--port", dest="port", help="the service's port",
 183                       default=9082)
 184
 185     (options, args) = parser.parse_args()
 186     status = service_ctrl(options.command, options.port)
 187     sys.exit(status)
 188
 189
 190 if __name__ == '__main__':
 191     main()