CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In
rapid7

CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

GitHub Repository: rapid7/metasploit-framework
Path: blob/master/lib/rex/job_container.rb
Views: 1904
1
# -*- coding: binary -*-
2
module Rex
3
4
###
5
#
6
# This class contains zero or more abstract jobs that can be enumerated and
7
# stopped in a generic fashion. This is used to provide a mechanism for
8
# keeping track of arbitrary contexts that may or may not require a dedicated
9
# thread.
10
#
11
###
12
class JobContainer < Hash
13
14
def initialize
15
self.job_id_pool = 0
16
end
17
18
#
19
# Adds an already running task as a symbolic job to the container.
20
#
21
def add_job(name, ctx, run_proc, clean_proc)
22
real_name = name
23
count = 0
24
jid = job_id_pool
25
26
self.job_id_pool += 1
27
28
# If we were not supplied with a job name, pick one from the hat
29
if (real_name == nil)
30
real_name = '#' + jid.to_s
31
end
32
33
# Find a unique job name
34
while (j = self[real_name])
35
real_name = name + " #{count}"
36
count += 1
37
end
38
39
j = Job.new(self, jid, real_name, ctx, run_proc, clean_proc)
40
41
self[jid.to_s] = j
42
end
43
44
#
45
# Starts a job using the supplied name and run/clean procedures.
46
#
47
def start_job(name, ctx, run_proc, clean_proc = nil)
48
j = add_job(name, ctx, run_proc, clean_proc)
49
j.start
50
51
j.jid
52
end
53
54
#
55
# Starts a background job that doesn't call the cleanup routine or run
56
# the run_proc in its own thread. Rather, the run_proc is called
57
# immediately and the clean_proc is never called until the job is removed
58
# from the job container.
59
#
60
def start_bg_job(name, ctx, run_proc, clean_proc = nil, async = true)
61
j = add_job(name, ctx, run_proc, clean_proc)
62
j.start(async)
63
64
j.jid
65
end
66
67
#
68
# Stops the job with the supplied name and forces it to cleanup. Stopping
69
# the job also leads to its removal.
70
#
71
def stop_job(jid)
72
if (j = self[jid.to_s])
73
j.stop
74
75
remove_job(j)
76
end
77
end
78
79
#
80
# Removes a job that was previously running. This is typically called when
81
# a job completes its task.
82
#
83
def remove_job(inst)
84
self.delete(inst.jid.to_s)
85
end
86
87
#
88
# Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
89
# "can't add a new key into hash during iteration"
90
#
91
def each(&block)
92
list = []
93
self.keys.sort.each do |sidx|
94
list << [sidx, self[sidx]]
95
end
96
list.each(&block)
97
end
98
99
protected
100
101
attr_accessor :job_id_pool # :nodoc:
102
103
end
104
105
end
106
107
108