Query & save runtime info of PD (#6)

* scripts: add helper func to wrap url request

* scripts: add basic functions to call `pd-ctl` apis

* pdctl: add argument of pd host/port

* pdctl: read runtime info from `pd-ctl`

* remove unused import

insight.py

@@ -29,6 +29,7 @@
 from measurement.files import configfiles
 from measurement.files import fileutils
 from measurement.files import logfiles
+from measurement.tidb import pdctl
 
 
 class Insight():
@@ -39,6 +40,7 @@ class Insight():
     insight_perf = None
     insight_logfiles = None
     insight_configfiles = None
+    insight_pdctl = None
 
     def __init__(self, outdir=None):
         self.full_outdir = fileutils.create_dir(self.outdir)
@@ -164,6 +166,10 @@ def save_configs(self, args):
         self.insight_configfiles.save_tidb_configs(
             proc_cmdline=proc_cmdline, outputdir=self.outdir)
 
+    def read_pdctl(self, args):
+        self.insight_pdctl = pdctl.PDCtl(host=args.pd_host, port=args.pd_port)
+        self.insight_pdctl.save_info(self.full_outdir)
+
 
 if __name__ == "__main__":
     if not util.is_root_privilege():
@@ -189,3 +195,5 @@ def save_configs(self, args):
     insight.save_logfiles(args)
     # save config files
     insight.save_configs(args)
+    # read and save `pd-ctl` info
+    insight.read_pdctl(args)

measurement/tidb/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

measurement/tidb/pdctl.py

@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+# Collect infomation with PD Controller
+
+import os
+
+from measurement import util
+from measurement.files import fileutils
+
+
+class PDCtl():
+    # default output dir name
+    pdctl_dir = "pdctl"
+
+    # default to localhost
+    pd_host = "localhost"
+    pd_port = 2379
+
+    # The `pdctl` API base URI
+    base_uri = "/pd/api"
+    base_url = ""
+
+    # The `pdctl` API version
+    api_ver = "1"
+    api_path = "/v%s" % api_ver
+
+    # The `pdctl` API URIs and relevant readable names, for full list of API
+    # routes, see: https://github.com/pingcap/pd/blob/master/server/api/router.go
+    # NOTE: All requests are GET
+    api_map = {
+        "config": "/config",
+        "operators": "/operators",
+        "schedulers": "/schedulers",
+        "labels": "/labels",
+        "hotspot": "/hotspot/stores",
+        "regions": "/regions",
+        "regionstats": "/stats/region",
+        "status": "/status",
+        "members": "/members",
+    }
+
+    # primary info of PD
+    pd_health_uri = "/health"
+    pd_diagnose_uri = "/diagnose"
+
+    def __init__(self, host=None, port=None, api_ver=None):
+        if host:
+            self.pd_host = host
+        if port:
+            self.pd_port = port
+        if api_ver:
+            self.api_ver = api_ver
+        self.base_url = "http://%s:%s%s%s" % (
+            self.pd_host, self.pd_port, self.base_uri, self.api_path)
+
+    def read_health(self):
+        url = "http://%s:%s/pd%s" % (self.pd_host,
+                                     self.pd_port, self.pd_health_uri)
+        return util.read_url(url)
+
+    def read_diagnose(self):
+        url = "http://%s:%s/pd%s" % (self.pd_host,
+                                     self.pd_port, self.pd_diagnose_uri)
+        return util.read_url(url)
+
+    def read_runtime_info(self):
+        def build_url(uri):
+            return "%s/%s" % (self.base_url, uri)
+
+        runtime_info = {}
+        for key, uri in self.api_map.items():
+            runtime_info[key] = util.read_url(build_url(uri))
+        return runtime_info
+
+    def save_info(self, basedir=None):
+        full_outputdir = fileutils.build_full_output_dir(
+            basedir=basedir, subdir=self.pdctl_dir)
+        fileutils.write_file(os.path.join(
+            full_outputdir, "%s-health.json" % self.pd_host), self.read_health())
+        fileutils.write_file(os.path.join(
+            full_outputdir, "%s-diagnose.json" % self.pd_host), self.read_diagnose())
+
+        for key, info in self.read_runtime_info().items():
+            fileutils.write_file(os.path.join(
+                full_outputdir, "%s-%s.json" % (self.pd_host, key)), info)

measurement/util.py

@@ -6,6 +6,14 @@
 import os
 
 from subprocess import Popen, PIPE
+try:
+    # For Python 2
+    import urllib2 as urlreq
+    from urllib2 import HTTPError, URLError
+except ImportError:
+    # For Python 3
+    import urllib.request as urlreq
+    from urllib.error import HTTPError, URLError
 
 
 def is_root_privilege():
@@ -72,6 +80,10 @@ def parse_insight_opts():
                         help="Enable to include system log in output, will be ignored if -l/--log is not set. This may significantly increase output size.")
     parser.add_argument("--config-file", action="store_true", default=False,
                         help="Enable to include various config files in output, disabled by default.")
+    parser.add_argument("--pd-host", action="store", default=None,
+                        help="The host of PD server. Default to localhost.")
+    parser.add_argument("--pd-port", type=int, action="store", default=None,
+                        help="The port of PD API service, default to 2379.")
 
     return parser.parse_args()
 
@@ -83,3 +95,17 @@ def get_init_type():
         logging.warning("Unable to detect init type, am I running with root?")
         return None
     return init_exec.split("/")[-1]
+
+
+def read_url(url, data=None):
+    if not url or url == "":
+        return None
+
+    try:
+        response = urlreq.urlopen(url, data)
+        return response.read()
+    except HTTPError as e:
+        return e.read()
+    except URLError as e:
+        logging.critical("Reading URL %s error: %s" % (url, e))
+        return None