damo_record: Move target pids polling to _damo_records

Signed-off-by: SeongJae Park <sj38.park@gmail.com>

2 files changed, 59 insertions, 59 deletions

diff --git a/_damo_records.py b/_damo_records.py
index f9dd3ee3..ee971240 100644
--- a/_damo_records.py
+++ b/_damo_records.py
@@ -577,6 +577,63 @@ def record_mem_footprint(kdamonds, snapshots):
 
 # for recording
 
+def pid_running(pid):
+    '''pid should be string'''
+    try:
+        subprocess.check_output(['ps', '--pid', pid])
+        return True
+    except:
+        return False
+
+def all_targets_terminated(targets):
+    for target in targets:
+        if pid_running('%s' % target.pid):
+            return False
+    return True
+
+def poll_target_pids(kdamonds, add_childs):
+    '''Return if polling should continued'''
+    current_targets = kdamonds[0].contexts[0].targets
+    if all_targets_terminated(current_targets):
+        return False
+    if not add_childs:
+        return True
+
+    for target in current_targets:
+        if target.pid == None:
+            continue
+        try:
+            childs_pids = subprocess.check_output(
+                    ['ps', '--ppid', '%s' % target.pid, '-o', 'pid=']
+                    ).decode().split()
+        except:
+            childs_pids = []
+        if len(childs_pids) == 0:
+            continue
+
+        # TODO: Commit all at once, out of this loop
+        new_targets = []
+        for child_pid in childs_pids:
+            # skip the child if already in the targets
+            if child_pid in ['%s' % t.pid for t in current_targets]:
+                continue
+            # remove already terminated targets, since committing already
+            # terminated targets to DAMON fails
+            new_targets = [target for target in current_targets
+                    if pid_running('%s' % target.pid)]
+            new_targets.append(_damon.DamonTarget(pid=child_pid, regions=[]))
+        if new_targets == []:
+            continue
+
+        # commit the new set of targets
+        kdamonds[0].contexts[0].targets = new_targets
+        err = _damon.commit(kdamonds)
+        if err != None:
+            # this might be not a problem; some of processes might
+            # finished meanwhile
+            return False
+    return True
+
 class RecordingHandle:
     file_path = None
     file_format = None
diff --git a/damo_record.py b/damo_record.py
index a174ed2b..046f33a6 100644
--- a/damo_record.py
+++ b/damo_record.py
@@ -71,64 +71,6 @@ def handle_args(args):
         print(err)
         exit(-3)
 
-def pid_running(pid):
-    '''pid should be string'''
-    try:
-        subprocess.check_output(['ps', '--pid', pid])
-        return True
-    except:
-        return False
-
-def all_targets_terminated(targets):
-    for target in targets:
-        if pid_running('%s' % target.pid):
-            return False
-    return True
-
-def poll_target_pids(kdamonds, add_childs):
-    '''Return if polling should continued'''
-    current_targets = kdamonds[0].contexts[0].targets
-    if all_targets_terminated(current_targets):
-        return False
-    if not add_childs:
-        return True
-
-    for target in current_targets:
-        if target.pid == None:
-            continue
-        try:
-            childs_pids = subprocess.check_output(
-                    ['ps', '--ppid', '%s' % target.pid, '-o', 'pid=']
-                    ).decode().split()
-        except:
-            childs_pids = []
-        if len(childs_pids) == 0:
-            continue
-
-        # TODO: Commit all at once, out of this loop
-        new_targets = []
-        for child_pid in childs_pids:
-            # skip the child if already in the targets
-            if child_pid in ['%s' % t.pid for t in current_targets]:
-                continue
-            # remove already terminated targets, since committing already
-            # terminated targets to DAMON fails
-            new_targets = [target for target in current_targets
-                    if pid_running('%s' % target.pid)]
-            new_targets.append(_damon.DamonTarget(pid=child_pid, regions=[]))
-        if new_targets == []:
-            continue
-
-        # commit the new set of targets
-        kdamonds[0].contexts[0].targets = new_targets
-        err = _damon.commit(kdamonds)
-        if err != None:
-            # this might be not a problem; some of processes might
-            # finished meanwhile
-            print('adding child as target failed (%s)' % err)
-            cleanup_exit(1)
-    return True
-
 def main(args):
     global data_for_cleanup
 
@@ -178,7 +120,8 @@ def main(args):
     print('Press Ctrl+C to stop')
 
     if _damon_args.self_started_target(args):
-        while poll_target_pids(kdamonds, args.include_child_tasks):
+        while _damo_records.poll_target_pids(kdamonds,
+                                             args.include_child_tasks):
             if args.footprint:
                 _damo_records.record_mem_footprint(
                         kdamonds, footprint_snapshots)