Issue 2239 - Make logprocessor close files and wait for subprocesses

sitescripts/stats/bin/logprocessor.py

 # coding: utf-8
 
 # This file is part of the Adblock Plus web scripts,
 # Copyright (C) 2006-2015 Eyeo GmbH
 #
 # Adblock Plus is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 3 as
 # published by the Free Software Foundation.
 #
 # Adblock Plus is distributed in the hope that it will be useful,
@@ skipping 24 matching lines @@
 import traceback
 import urllib
 import urlparse
 
 import sitescripts.stats.common as common
 from sitescripts.utils import get_config, setupStderr
 
 log_regexp = None
 gecko_apps = None
 
-def open_stats_file(path):
-  parseresult = urlparse.urlparse(path)
-  if parseresult.scheme == "ssh" and parseresult.username and parseresult.hostname and parseresult.path:
-    command = [
-      "ssh", "-q", "-o", "NumberOfPasswordPrompts 0", "-T", "-k",
-      "-l", parseresult.username,
-      parseresult.hostname,
-      parseresult.path.lstrip("/")
-    ]
-    if parseresult.port:
-      command[1:1] = ["-P", str(parseresult.port)]
-    result = subprocess.Popen(command, stdout=subprocess.PIPE).stdout
-  elif parseresult.scheme in ("http", "https"):
-    result = urllib.urlopen(path)
-  elif os.path.exists(path):
-    result = open(path, "rb")
-  else:
-    raise IOError("Path '%s' not recognized" % path)
+class StatsFile:
+  def __init__(self, path):
+    self._inner_file = None
+    self._processes = []
 
-  if path.endswith(".gz"):
-    # Built-in gzip module doesn't support streaming (fixed in Python 3.2)
-    result = subprocess.Popen(["gzip", "-cd"], stdin=result, stdout=subprocess.PIPE).stdout
-  return result
+    parseresult = urlparse.urlparse(path)
+    if parseresult.scheme == "ssh" and parseresult.username and parseresult.hostname and parseresult.path:
+      command = [
+        "ssh", "-q", "-o", "NumberOfPasswordPrompts 0", "-T", "-k",
+        "-l", parseresult.username,
+        parseresult.hostname,
+        parseresult.path.lstrip("/")
+      ]
+      if parseresult.port:
+        command[1:1] = ["-P", str(parseresult.port)]
+      ssh_process = subprocess.Popen(command, stdout=subprocess.PIPE)
+      self._processes.push(ssh_process)
+      self._file = ssh_process.stdout
+    elif parseresult.scheme in ("http", "https"):
+      self._file = urllib.urlopen(path)
+    elif os.path.exists(path):
+      self._file = open(path, "rb")
+    else:
+      raise IOError("Path '%s' not recognized" % path)
+
+    if path.endswith(".gz"):
+      # Built-in gzip module doesn't support streaming (fixed in Python 3.2)
+      gzip_process = subprocess.Popen(["gzip", "-cd"], stdin=self._file, stdout=subprocess.PIPE)
+      self._processes.append(gzip_process)
+      self._file, self._inner_file = gzip_process.stdout, self._file
+
+  def __getattr__(self, name):
+    return getattr(self._file, name)
+
+  def close(self):
+    self._file.close()
+    if self._inner_file:
+      self._inner_file.close()
+    for process in self._processes:
+      process.wait()
 
 def get_stats_files():
   config = get_config()
 
   prefix = "mirror_"
   options = filter(lambda o: o.startswith(prefix), config.options("stats"))
   for option in options:
     if config.has_option("stats", option):
       value = config.get("stats", option)
       if " " in value:
@@ skipping 423 matching lines @@
 
       with codecs.open(path, "wb", encoding="utf-8") as fileobj:
         json.dump(existing, fileobj, indent=2, sort_keys=True)
 
 def parse_source(factor, lock, (mirror_name, server_type, log_file)):
   try:
     geo = pygeoip.GeoIP(get_config().get("stats", "geoip_db"), pygeoip.MEMORY_CACHE)
     geov6 = pygeoip.GeoIP(get_config().get("stats", "geoipv6_db"), pygeoip.MEMORY_CACHE)
 
     ignored = set()
-    fileobj = open_stats_file(log_file)
+    fileobj = StatsFile(log_file)
     try:
       data = parse_fileobj(mirror_name, fileobj, geo, geov6, ignored)
     finally:
       fileobj.close()
 
     lock.acquire()
     try:
       save_stats(server_type, data, factor)
     finally:
       lock.release()
@@ skipping 25 matching lines @@
   parser.add_argument("mirror_name", nargs="?", help="Name of the mirror server that the file belongs to")
   parser.add_argument("server_type", nargs="?", help="Server type like download, update or subscription")
   parser.add_argument("log_file", nargs="?", help="Log file path, can be a local file path, http:// or ssh:// URL")
   args = parser.parse_args()
 
   if args.mirror_name and args.server_type and args.log_file:
     sources = [(args.mirror_name, args.server_type, args.log_file)]
   else:
     sources = get_stats_files()
   parse_sources(sources, args.factor, args.verbose)