Merge pull request #278 from shreyasp/bench/rate_limit_conn

[Minor] Nginx rate limiting and Improvements to Bench Installer

bench/config/nginx.py

@@ -1,4 +1,4 @@
-import os, json, click, random, string
+import os, json, click, random, string, hashlib
 from bench.utils import get_sites, get_bench_name, exec_cmd
 
 def make_nginx_conf(bench_path, yes=False):
@@ -12,13 +12,18 @@ def make_nginx_conf(bench_path, yes=False):
 	config = get_config(bench_path)
 	sites = prepare_sites(config, bench_path)
 
+	bench_name = get_bench_name(bench_path)
+	bench_name_hash = hashlib.sha256(bench_name).hexdigest()[:16]
+
 	nginx_conf = template.render(**{
 		"sites_path": sites_path,
 		"http_timeout": config.get("http_timeout"),
 		"sites": sites,
 		"webserver_port": config.get('webserver_port'),
 		"socketio_port": config.get('socketio_port'),
-		"bench_name": get_bench_name(bench_path),
+		"bench_name": bench_name,
+		"bench_name_hash": bench_name_hash,
+		"limit_conn_shared_memory": get_limit_conn_shared_memory(),
 		"error_pages": get_error_pages(),
 
 		# for nginx map variable
@@ -156,3 +161,10 @@ def get_error_pages():
 	return {
 		502: os.path.join(templates, '502.html')
 	}
+
+def get_limit_conn_shared_memory():
+	"""Allocate 2 percent of total virtual memory as shared memory for nginx limit_conn_zone"""
+	import psutil
+	total_vm = (psutil.virtual_memory().total) / (1024 * 1024) # in MB
+
+	return int(0.02 * total_vm)

bench/config/production_setup.py

@@ -127,5 +127,9 @@ def reload_supervisor():
 		pass
 
 def reload_nginx():
-	subprocess.check_output(['sudo', find_executable('nginx'), '-t'])
+	try:
+		subprocess.check_output(['sudo', find_executable('nginx'), '-t'])
+	except:
+		raise
+
 	service('nginx', 'reload')

bench/config/templates/nginx.conf

@@ -21,6 +21,8 @@ server {
 
 	root {{ sites_path }};
 
+	limit_conn per_host_{{ bench_name_hash }} 8;
+
 	{% if ssl_certificate and ssl_certificate_key %}
 	ssl on;
 	ssl_certificate      {{ ssl_certificate }};
@@ -139,6 +141,8 @@ upstream {{ bench_name}}-socketio-server {
     server 127.0.0.1:{{ socketio_port or 3000 }} fail_timeout=0;
 }
 
+limit_conn_zone $host zone=per_host_{{ bench_name_hash }}:{{ limit_conn_shared_memory }}m;
+
 # setup maps
 {%- set site_name_variable="$host" %}
 {% if sites.domain_map -%}

playbooks/install.py

@@ -5,6 +5,8 @@ from distutils.spawn import find_executable
 tmp_bench_repo = '/tmp/.bench'
 
 def install_bench(args):
+	check_brew_installed()
+
 	# pre-requisites for bench repo cloning
 	install_package('curl')
 	install_package('wget')
@@ -146,6 +148,20 @@ def install_package(package):
 	if not success:
 		could_not_install(package)
 
+def check_brew_installed():
+	if 'Darwin' not in os.uname():
+		return
+
+	brew_exec = find_executable('brew')
+
+	if not brew_exec:
+		raise Exception('''
+		Please install brew package manager before proceeding with bench setup. Please run following
+		to install brew package manager on your machine,
+
+		/usr/bin/ruby -e "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install)"
+		''')
+
 def clone_bench_repo(args):
 	'''Clones the bench repository in the user folder'''
 	if os.path.exists(tmp_bench_repo):

requirements.txt

@@ -3,6 +3,7 @@ jinja2
 virtualenv
 requests
 honcho
+psutil
 python-crontab
 semantic_version
 GitPython==0.3.2.rc1