diff --git a/bin/zk-stats-daemon b/bin/zk-stats-daemon index 9239e29..f294550 100755 --- a/bin/zk-stats-daemon +++ b/bin/zk-stats-daemon @@ -11,6 +11,7 @@ if os.getenv("ZKTRAFFIC_SOURCE") is not None: sys.path.insert(0, ".") from zktraffic.endpoints.stats_server import StatsServer +from zktraffic.base.process import ProcessOptions from twitter.common import app, log from twitter.common.http import HttpServer @@ -57,6 +58,12 @@ app.add_option("--niceness", type=int, default=0, help="set the niceness") +app.add_option("--set-cpu-affinity", + dest="cpu_affinity", + metavar="CPU#[,CPU#]", + type=str, + default=None, + help="A comma-separated list of CPU cores to pin this process to") app.add_option("--max-queued-requests", type=int, default=400000, @@ -88,8 +95,13 @@ def main(_, opts): signal.signal(signal.SIGINT, signal.SIG_DFL) + process = ProcessOptions() + if opts.niceness >= 0: - os.nice(opts.niceness) + process.set_niceness(opts.niceness) + + if opts.cpu_affinity: + process.set_cpu_affinity(opts.cpu_affinity) server = Server() server.mount_routes(stats) diff --git a/requirements.txt b/requirements.txt index 216f105..17d87bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ dpkt-fix nose -twitter.common.log +psutil>=2.1.0 scapy==2.2.0-dev +twitter.common.log diff --git a/setup.py b/setup.py index 4b24310..4f03e4f 100644 --- a/setup.py +++ b/setup.py @@ -63,18 +63,20 @@ def get_version(): install_requires=[ 'ansicolors', 'dpkt-fix', + 'psutil>=2.1.0', + 'scapy==2.2.0-dev', 'twitter.common.app', 'twitter.common.collections', 'twitter.common.exceptions', 'twitter.common.http', 'twitter.common.log', - 'scapy==2.2.0-dev', ], tests_require=[ 'dpkt-fix', 'nose', - 'twitter.common.log', + 'psutil>=2.1.0', 'scapy==2.2.0-dev', + 'twitter.common.log', ], extras_require={ 'test': [ diff --git a/zktraffic/base/process.py b/zktraffic/base/process.py new file mode 100644 index 0000000..f89f380 --- /dev/null +++ b/zktraffic/base/process.py @@ -0,0 +1,78 @@ +# ================================================================================================== +# Copyright 2014 Twitter, Inc. +# -------------------------------------------------------------------------------------------------- +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this work except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file, or at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ================================================================================================== + + +""" Process helpers used to improve zktraffic operation """ + +from twitter.common import log + +import psutil +from psutil import AccessDenied, NoSuchProcess + + +class ProcessOptions(object): + + def __init__(self): + self.process = psutil.Process() + + def set_cpu_affinity(self, cpu_affinity_csv): + """ + Set CPU affinity for this process + :param cpu_affinity_csv: A comma-separated string representing CPU cores + """ + try: + cpu_list = self.parse_cpu_affinity(cpu_affinity_csv) + self.process.cpu_affinity(cpu_list) + except (OSError, ValueError) as e: + log.warn('unable to set cpu affinity: {}, on process: {}'.format(cpu_affinity_csv, e)) + except AttributeError: + log.warn('cpu affinity is not available on your platform') + + def get_cpu_affinity(self): + """ + Get CPU affinity of this process + :return: a list() of CPU cores this processes is pinned to + """ + return self.process.cpu_affinity() + + def set_niceness(self, nice_level): + """ + Set the nice level of this process + :param nice_level: the nice level to set + """ + try: + # TODO (phobos182): double check that psutil does not allow negative nice values + if not 0 <= nice_level <= 20: + raise ValueError('nice level must be between 0 and 20') + self.process.nice(nice_level) + except(EnvironmentError, ValueError, AccessDenied, NoSuchProcess) as e: + log.warn('unable to set nice level on process: {}'.format(e)) + + def get_niceness(self): + """ + Get nice level of this process + :return: an int() representing the nice level of this process + """ + return self.process.nice() + + @staticmethod + def parse_cpu_affinity(cpu_affinity_csv): + """ + Static method to parse a csv string string into a list of integers + :param cpu_affinity_csv: a CSV of cpu cores to parse + :return: a list() of integers representing CPU cores + """ + return [int(_) for _ in cpu_affinity_csv.split(',')] \ No newline at end of file diff --git a/zktraffic/tests/test_process_options.py b/zktraffic/tests/test_process_options.py new file mode 100644 index 0000000..622975a --- /dev/null +++ b/zktraffic/tests/test_process_options.py @@ -0,0 +1,61 @@ +# ================================================================================================== +# Copyright 2014 Twitter, Inc. +# -------------------------------------------------------------------------------------------------- +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this work except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file, or at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ================================================================================================== + + +import mock +import os + +from zktraffic.base.process import ProcessOptions + +import psutil + + +proc = ProcessOptions() + + +def test_niceness(): + """ + Test CPU niceness calls + """ + proc.set_niceness(15) + assert proc.get_niceness() == 15 + + +def test_cpu_affinity_parsing(): + """ + Test CPU affinity list parsing + """ + assert ProcessOptions.parse_cpu_affinity('0,1,2,3,4,5') == [0, 1, 2, 3, 4, 5] + assert ProcessOptions.parse_cpu_affinity('2') == [2] + + +def test_cpu_affinity(): + """ + Test CPU affinity setting + """ + def mock_cpu_affinity_handler(self, *args, **kwargs): + return [0, 1] + + # if running in TravisCI, mock the cpu_affinity() method call + if os.environ.get('TRAVIS'): + with mock.patch.object(psutil.Process, 'cpu_affinity', create=True, new=mock_cpu_affinity_handler): + proc.set_cpu_affinity('0,1') + assert proc.get_cpu_affinity() == [0, 1] + else: + proc.set_cpu_affinity('0,1') + assert proc.get_cpu_affinity() == [0, 1] + proc.set_cpu_affinity('1') + assert proc.get_cpu_affinity() == [1] \ No newline at end of file