00001
00002
00003 import rospy
00004 import time
00005 import subprocess
00006 import signal
00007 import os
00008 import sys
00009 import argparse
00010 import re
00011 import shutil
00012 from jsk_topic_tools.master_util import isMasterAlive
00013
00014
00015 def runROSBag(topics, size, save_dir):
00016 """
00017 run rosbag and return Popen object
00018 """
00019 cmd = 'rosbag record --split'
00020 formatted_topics = [t for t in topics.split(' ') if t]
00021 return subprocess.Popen(cmd.split(' ')
00022 + formatted_topics
00023 + ['--size', size]
00024 + ['-o', save_dir + '/rosbag'])
00025
00026 def parseBagFile(bag):
00027
00028
00029 regexp = 'rosbag_(\d\d\d\d)-(\d\d)-(\d\d)-(\d\d)-(\d\d)-(\d\d)_([\d]+).bag'
00030 result = re.match(regexp, bag)
00031 if not result:
00032 return None
00033 else:
00034 return [result.group(f) for f in range(1, 8)]
00035
00036 def mkdirForBag(root_dir, bag):
00037
00038 parse = parseBagFile(bag)
00039 YYYY = parse[0]
00040 MM = parse[1]
00041 DD = parse[2]
00042 HH = parse[3]
00043 directory = os.path.join(root_dir, YYYY, MM, DD, HH)
00044 if not os.path.exists(directory):
00045 os.makedirs(directory)
00046 return directory
00047
00048 def moveBagFiles(save_dir, bags):
00049 for bag in bags:
00050 move_dir = mkdirForBag(save_dir, bag)
00051 from_file = os.path.join(save_dir, bag)
00052 to_file = os.path.join(move_dir, bag)
00053 print 'moving file %s -> %s' % (from_file, to_file)
00054 shutil.move(from_file, to_file)
00055
00056 def watchFileSystem(save_dir, max_size):
00057 files = os.listdir(save_dir)
00058 target_bags = [f for f in files
00059 if f.startswith('rosbag')
00060 and f.endswith('.bag')]
00061 moveBagFiles(save_dir, target_bags)
00062
00063 checkDirectorySize(save_dir, int(max_size))
00064
00065 def getDirectorySize(start_path = '.'):
00066 "size unit is Mega bytes"
00067 total_size = 0
00068 for dirpath, dirnames, filenames in os.walk(start_path):
00069 for f in filenames:
00070 fp = os.path.join(dirpath, f)
00071 try:
00072 total_size += os.path.getsize(fp)
00073 except:
00074 pass
00075 return total_size / 1000.0 / 1000.0
00076
00077 def keyFuncToSortBag(bag):
00078 parse = parseBagFile(os.path.basename(bag))
00079 return int(reduce(lambda x, y: x + y, parse))
00080
00081 def listBagsSortedByDate(save_dir):
00082 bags = []
00083 for dirpath, dirnames, filenames in os.walk(save_dir):
00084 for f in filenames:
00085 if f.startswith('rosbag') and f.endswith('.bag'):
00086 bags.append(os.path.join(dirpath, f))
00087 return sorted(bags, key=keyFuncToSortBag)
00088
00089 def removeOldFiles(save_dir, max_size, current_size):
00090 files = listBagsSortedByDate(save_dir)
00091 remove_size = current_size - max_size
00092 for f in files:
00093 the_size = os.path.getsize(f)
00094 print 'removing %s (%d)' % (f, the_size / 1000 / 1000)
00095 os.remove(f)
00096 remove_size = remove_size - the_size / 1000.0 / 1000.0
00097 if remove_size < 0:
00098 return
00099
00100 def checkDirectorySize(save_dir, max_size):
00101 size = getDirectorySize(save_dir)
00102 print 'current directory size is %fM (max is %dM)' % (size, int(max_size))
00103
00104 if size > max_size:
00105 removeOldFiles(save_dir, max_size, size)
00106
00107 g_rosbag_process = False
00108
00109 def restartROSBag(topics, size, save_dir):
00110 global g_rosbag_process
00111 print 'Running rosbag...'
00112 g_rosbag_process = runROSBag(topics, size, save_dir)
00113
00114 def killChildProcesses(ppid):
00115 output = subprocess.check_output(['ps', '--ppid=' + str(ppid), '--no-headers'])
00116 for process_line in output.split('\n'):
00117 strip_process_line = process_line.strip()
00118 if strip_process_line:
00119 pid = strip_process_line.split(' ')[0]
00120 name = strip_process_line.split(' ')[-1]
00121 print 'killing %s' % (name)
00122 os.kill(int(pid), signal.SIGINT)
00123
00124 def killROSBag():
00125 global g_rosbag_process
00126 if g_rosbag_process:
00127 print 'Killing rosbag ...'
00128 rosbag_pid = g_rosbag_process.pid
00129 killChildProcesses(rosbag_pid)
00130 g_rosbag_process.send_signal(subprocess.signal.SIGINT)
00131
00132 def main(topics, size, save_dir, max_size, rate = 1):
00133 if not os.path.exists(save_dir):
00134 os.makedirs(save_dir)
00135 previous_master_state = None
00136 try:
00137 while True:
00138 master_state = isMasterAlive()
00139 if not master_state and previous_master_state:
00140 print "kill rosbag"
00141 killROSBag()
00142 elif master_state and not previous_master_state:
00143 print "restart rosbag"
00144 restartROSBag(topics, size, save_dir)
00145 watchFileSystem(save_dir, max_size)
00146 previous_master_state = master_state
00147 time.sleep(1.0 / rate)
00148 except:
00149 time.sleep(1)
00150 watchFileSystem(save_dir, max_size)
00151
00152
00153
00154 if __name__ == "__main__":
00155 parser = argparse.ArgumentParser(description='rosbag record regardless of rosmaster status')
00156 parser.add_argument('--topics', help="topics to record", required=True)
00157 parser.add_argument('--size', help="size of each rosbag", required=True)
00158 parser.add_argument('--save-dir', help="directory to store rosbag", required=True)
00159 parser.add_argument('--max-size', help="maximum size of rosbags in save_dir", required=True, type=int)
00160 args = parser.parse_args()
00161 main(args.topics, args.size, args.save_dir, args.max_size)