Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-09-28 23:48:30

0001 #!/usr/bin/env python3
0002 
0003 import argparse
0004 import subprocess
0005 import select, signal, time, errno
0006 import sys, os
0007 import zlib
0008 
0009 def log(s):
0010     sys.stderr.write("m: " + s + "\n");
0011     sys.stderr.flush()
0012 
0013 class GZipLog(object):
0014     def __init__(self, log_file):
0015         self.file = None
0016 
0017         self.last_flush = time.time()
0018         self.zstream = zlib.compressobj(6, zlib.DEFLATED, zlib.MAX_WBITS | 16)
0019 
0020         # 'gzip' block will be flushed if we reach max bytes or the timeout
0021         self.block = []
0022         self.block_size = 0
0023         self.block_size_max = 2*1024*1024 # bytes
0024         self.block_timeout = 15 # seconds
0025 
0026         self.file = open(log_file, "wb+")
0027         self.file_min_size = 2*1024*1024    # this amount of space we keep after truncating
0028         self.file_max_size = 16*1024*1024   # we truncate if file gets bigger
0029 
0030         self.file_truncate_pos = None
0031         self.file_truncate_state = None
0032 
0033     def write_block(self, data):
0034         self.file.write(self.zstream.compress( data.encode("utf-8") ))
0035         self.file.write(self.zstream.flush(zlib.Z_FULL_FLUSH))
0036 
0037     def flush_block(self):
0038         self.last_flush = time.time()
0039 
0040         data, size, = "".join(self.block), self.block_size
0041         self.block, self.block_size = [], 0
0042 
0043         if size == 0:
0044             # nothing to do
0045             return
0046 
0047         # check if we need to truncate the file
0048         # 'size' refers to uncompressed data, so the file be truncated at max_size -+ max_block_size
0049         # there is no way know compressed size without collapsing the zlib state
0050         if (self.file.tell() + size) > self.file_max_size:
0051             if self.file_truncate_pos is not None:
0052                 # tell it to cleanup
0053                 self.zstream.flush(zlib.Z_FINISH)
0054 
0055                 self.file.seek(self.file_truncate_pos, 0)
0056                 self.file.truncate()
0057                 self.zstream = self.file_truncate_state.copy()
0058 
0059                 self.write_block("\n\n--- file was cut at this point ---\n\n")
0060 
0061         # write it
0062         self.write_block(data)
0063 
0064         # check if we can save the ptr into truncate_position
0065         if self.file_truncate_pos is None:
0066             if self.file.tell() >= self.file_min_size:
0067                 self.file_truncate_pos = self.file.tell()
0068                 self.file_truncate_state = self.zstream.copy()
0069 
0070     def finish(self):
0071         self.flush_block()
0072         self.file.write(self.zstream.flush(zlib.Z_FINISH))
0073         self.file.close()
0074         self.file = None
0075 
0076     def try_flush(self):
0077         timeout = (time.time() - self.last_flush) > self.block_timeout
0078         large_block = self.block_size > self.block_size_max
0079 
0080         if timeout or large_block:
0081             self.flush_block()
0082 
0083     def write(self, bytes):
0084         if bytes:
0085             self.block.append(bytes)
0086             self.block_size += len(bytes)
0087 
0088         self.try_flush()
0089 
0090     def handle_timeout(self):
0091         self.try_flush()
0092 
0093 
0094 def capture(fd, args):
0095     log_handler = GZipLog(log_file=args.log)
0096 
0097     def sigusr1_handle(*kargs, **kwargs):
0098         log_handler._sigusr_interrupt = True
0099         if log_handler.file:
0100             log_handler.flush_block()
0101 
0102     signal.signal(signal.SIGUSR1, sigusr1_handle)
0103 
0104     try:
0105         while True:
0106             try:
0107                 log_handler._sigusr_interrupt = False
0108                 rlist, wlist, xlist = select.select([fd], [], [], 5)
0109             except select.error as e:
0110                 if e[0] != errno.EINTR: raise
0111                 if not log_handler._sigusr_interrupt: raise
0112 
0113                 continue
0114 
0115             if not rlist:
0116                 log_handler.handle_timeout()
0117                 continue
0118 
0119             bytes = os.read(fd, 4096)
0120             if (bytes == ''):
0121                 break
0122 
0123             log_handler.write(bytes)
0124             if not args.q:
0125                 sys.stdout.write(bytes)
0126     finally:
0127         log_handler.finish()
0128 
0129 if __name__ == "__main__":
0130     parser = argparse.ArgumentParser(description="Take all input and write a compressed log file.")
0131     parser.add_argument('-q', action='store_true', help="Don't write to stdout, just the log file.")
0132     parser.add_argument("log", type=str, help="Filename to write.", metavar="<logfile.gz>")
0133     args = parser.parse_args()
0134 
0135     if args.q:
0136         fd = sys.stdout.fileno()
0137         sys.stdout.close()
0138         os.close(fd)
0139 
0140     capture(sys.stdin.fileno(), args)