first commit

.gitignore

@@ -0,0 +1,3 @@
+data/*
+*.pyc
+helpers/*.pyc

README.md

@@ -0,0 +1,16 @@
+# PostgreSQL HA with etcd
+
+To get started, do the following from different terminals:
+
+```
+> etcd --data-dir=data/etcd
+> run.py postgresql0.yml
+> run.py postgresql1.yml
+```
+
+From there, you will see a high-availability cluster start up. Test
+different settings in the YAML files to see how behavior changes.  Kill
+some of the different components to see how the system behaves.
+
+Cheers,
+Chris

helpers/errors.py

@@ -0,0 +1,13 @@
+class CurrentLeaderError(Exception):
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return repr(self.value)
+
+class HealthiestMemberError(Exception):
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return repr(self.value)

helpers/etcd.py

@@ -0,0 +1,102 @@
+import urllib2, json, os, time
+from urllib import urlencode
+import helpers.errors
+
+class Etcd:
+    def __init__(self, config):
+        self.scope = config["scope"]
+        self.host = config["host"]
+        self.ttl = config["ttl"]
+
+    def get_client_path(self, path, max_attempts = 1):
+        attempts = 0
+        response = None
+
+        while True:
+            try:
+                response = urllib2.urlopen(self.client_url(path)).read()
+                break
+            except (urllib2.HTTPError, urllib2.URLError) as e:
+                attempts += 1
+                if attempts < max_attempts:
+                    print "Failed to return %s, trying again. (%s of %s)" % (path, attempts, max_attempts)
+                    time.sleep(3)
+                else:
+                    raise e
+        try:
+            return json.loads(response)
+        except ValueError:
+            return response
+
+    def put_client_path(self, path, data):
+        opener = urllib2.build_opener(urllib2.HTTPHandler)
+        request = urllib2.Request(self.client_url(path), data=urlencode(data).replace("false", "False"))
+        request.get_method = lambda: 'PUT'
+        opener.open(request)
+
+    def client_url(self, path):
+        return "http://%s/v2/keys/service/%s%s" % (self.host, self.scope, path)
+
+    def xlog_position(member):
+        try:
+            return self.get_client_path("/service/postgresql/xlog-position/%s" % member)["node"]["value"]
+        except urllib2.HTTPError:
+            return None
+
+    def current_leader(self):
+        try:
+            hostname = self.get_client_path("/leader")["node"]["value"]
+            address = self.get_client_path("/members/%s" % hostname)["node"]["value"]
+
+            return {"hostname": hostname, "address": address}
+        except urllib2.HTTPError as e:
+            if e.code == 404:
+                return None
+            raise helpers.errors.CurrentLeaderError("Etcd is not responding properly")
+
+    def touch_member(self, member, connection_string):
+        self.put_client_path("/members/%s" % member, {"value": connection_string})
+
+    def take_leader(self, value):
+        return self.put_client_path("/leader", {"value": value, "ttl": self.ttl}) == None
+
+    def attempt_to_acquire_leader(self, value):
+        try:
+            return self.put_client_path("/leader", {"value": value, "ttl": self.ttl, "prevExist": False}) == None
+        except urllib2.HTTPError as e:
+            if e.code == 412:
+                print("Could not take out TTL lock: %s" % e)
+            return False
+
+    def update_leader(self, value):
+        try:
+            self.put_client_path("/leader", {"value": value, "ttl": self.ttl, "prevValue": value})
+        except urllib2.HTTPError:
+            print "Error updating TTL on ETCD for primary."
+            return False
+
+    def leader_unlocked(self):
+        try:
+            self.get_client_path("/leader")
+            return False
+        except urllib2.HTTPError as e:
+            if e.code == 404:
+                return True
+            return False
+        except ValueError as e:
+            return False
+
+    def am_i_leader(self, value):
+        #try:
+           reponse = self.get_client_path("/leader")
+           print "Lock owner: %s; I am %s" % (reponse["node"]["value"], value)
+           return reponse["node"]["value"] == value
+        #except Exception as e:
+            #return False
+
+
+    def race(self, path, value):
+        try:
+            return self.put_client_path(path, {"prevExist": False, "value": value}) == None
+        except urllib2.HTTPError:
+            return False

helpers/ha.py

@@ -0,0 +1,111 @@
+import sys, time, re, urllib2, json, psycopg2
+from base64 import b64decode
+
+import helpers.errors
+
+import inspect
+
+def lineno():
+    """Returns the current line number in our program."""
+    return inspect.currentframe().f_back.f_lineno
+
+class Ha:
+    def __init__(self, state_handler, etcd):
+        self.state_handler = state_handler
+        self.etcd = etcd
+
+    def acquire_lock(self):
+        return self.etcd.attempt_to_acquire_leader(self.state_handler.name)
+
+    def update_lock(self):
+        return self.etcd.update_leader(self.state_handler.name)
+
+    def is_unlocked(self):
+        return self.etcd.leader_unlocked()
+
+    def has_lock(self):
+        return self.etcd.am_i_leader(self.state_handler.name)
+
+    def fetch_current_leader(self):
+            return self.etcd.current_leader()
+
+    def run_cycle(self):
+        try:
+            print lineno()
+            if self.state_handler.is_healthy():
+                print lineno()
+                if self.is_unlocked():
+                    print lineno()
+                    if self.state_handler.is_healthiest_node():
+                        print lineno()
+                        if self.acquire_lock():
+                            print lineno()
+                            if not self.state_handler.is_leader():
+                                print lineno()
+                                self.state_handler.promote()
+                                return "promoted self to leader by acquiring session lock"
+
+                            print lineno()
+                            return "acquired session lock as a leader"
+                        else:
+                            print lineno()
+                            if self.state_handler.is_leader():
+                                print lineno()
+                                self.state_handler.demote(self.fetch_current_leader())
+                                return "demoted self due after trying and failing to obtain lock"
+                            else:
+                                print lineno()
+                                self.state_handler.follow_the_leader(self.fetch_current_leader())
+                                return "following new leader after trying and failing to obtain lock"
+                    else:
+                        print lineno()
+                        if self.state_handler.is_leader():
+                            print lineno()
+                            self.state_handler.demote(self.fetch_current_leader())
+                            return "demoting self because i am not the healthiest node"
+                        else:
+                            print lineno()
+                            self.state_handler.follow_the_leader(self.fetch_current_leader())
+                            return "following a different leader because i am not the healthiest node"
+
+                else:
+                    print lineno()
+                    if self.has_lock():
+                        print lineno()
+                        self.update_lock()
+
+                        if not self.state_handler.is_leader():
+                            print lineno()
+                            self.state_handler.promote()
+                            return "promoted self to leader because i had the session lock"
+                        else:
+                            print lineno()
+                            return "no action.  i am the leader with the lock"
+                    else:
+                        print lineno()
+                        print "does not have lock"
+                        if self.state_handler.is_leader():
+                            print lineno()
+                            self.state_handler.demote(self.fetch_current_leader())
+                            return "demoting self because i do not have the lock and i was a leader"
+                        else:
+                            print lineno()
+                            self.state_handler.follow_the_leader(self.fetch_current_leader())
+                            return "no action.  i am a secondary and i am following a leader"
+            else:
+                print lineno()
+                return "no action.  not healthy enough to do anything."
+        except helpers.errors.CurrentLeaderError:
+            print lineno()
+            print "failed to fetch current leader from etcd"
+        except psycopg2.OperationalError:
+            print lineno()
+            print "Error communicating with Postgresql.  Will try again."
+        except helpers.errors.HealthiestMemberError:
+            print lineno()
+            print "failed to determine healthiest member fromt etcd"
+
+    def run(self):
+        while True:
+            self.run_cycle()
+            time.sleep(10)

helpers/postgresql.py

@@ -0,0 +1,152 @@
+import os, psycopg2, re, time
+from urlparse import urlparse
+
+class Postgresql:
+
+    def __init__(self, config):
+        self.name = config["name"]
+        self.host, self.port = config["listen"].split(":")
+        self.data_dir = config["data_dir"]
+        self.replication = config["replication"]
+
+        self.config = config
+
+        self.cursor_holder = None
+        self.connection_string = "postgres://%s:%s@%s:%s/postgres" % (self.replication["username"], self.replication["password"], self.host, self.port)
+
+        self.conn = None
+
+    def cursor(self):
+        if self.cursor_holder == None:
+            self.conn = psycopg2.connect("postgres://%s:%s/postgres" % (self.host, self.port))
+            self.conn.autocommit = True
+            self.cursor_holder = self.conn.cursor()
+
+        return self.cursor_holder
+
+    def disconnect(self):
+        try:
+            self.conn.close()
+        except Exception as e:
+            print "Error disconnecting: %s" % e
+
+    def query(self, sql):
+        max_attempts = 0
+        while True:
+            try:
+                self.cursor().execute(sql)
+                break
+            except psycopg2.OperationalError as e:
+                if self.conn != None:
+                    self.disconnect()
+                self.cursor_holder = None
+                if max_attempts > 4:
+                    raise e
+                max_attempts += 1
+                time.sleep(5)
+        return self.cursor()
+
+    def data_directory_empty(self):
+        return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
+
+    def initialize(self):
+        if os.system("initdb -D %s" % self.data_dir) == 0:
+            self.write_pg_hba()
+
+            return True
+
+        return False
+
+    def sync_from_leader(self, leader):
+        leader = urlparse(leader["address"])
+
+        f = open("./pgpass", "w")
+        f.write("%(hostname)s:%(port)s:*:%(username)s:%(password)s\n" %
+                {"hostname": leader.hostname, "port": leader.port, "username": leader.username, "password": leader.password})
+        f.close()
+
+        os.system("chmod 600 pgpass")
+
+
+        return os.system("PGPASSFILE=pgpass pg_basebackup -R -D %(data_dir)s --host=%(host)s --port=%(port)s -U %(username)s" %
+                {"data_dir": self.data_dir, "host": leader.hostname, "port": leader.port, "username": leader.username}) == 0
+
+    def is_leader(self):
+        return not self.query("SELECT pg_is_in_recovery();").fetchone()[0]
+
+    def is_running(self):
+        return os.system("pg_ctl status -D %s" % self.data_dir) == 0
+
+    def start(self):
+        command_code = os.system("postgres -D %s %s &" % (self.data_dir, self.server_options()))
+        time.sleep(5)
+        return command_code != 0
+
+    def stop(self):
+        return os.system("pg_ctl stop -w -D %s -m fast -w" % self.data_dir) != 0
+
+    def reload(self):
+        return os.system("pg_ctl reload -w -D %s" % self.data_dir) == 0
+
+    def restart(self):
+        return os.system("pg_ctl restart -w -D %s -m fast" % self.data_dir) == 0
+
+    def server_options(self):
+        options = "-c listen_addresses=%s -c port=%s" % (self.host, self.port)
+        for setting, value in self.config["parameters"].iteritems():
+            options += " -c \"%s=%s\"" % (setting, value)
+        return options
+
+    def is_healthy(self):
+        if not self.is_running():
+            print "Postgresql is not running."
+            return False
+
+        return True
+
+    def is_healthiest_node(self):
+        return True
+
+    def replication_slot_name(self):
+        member = os.environ.get("MEMBER")
+        (member, _) = re.subn(r'[^a-z0-9]+', r'_', member)
+        return member
+
+    def write_pg_hba(self):
+        f = open("%s/pg_hba.conf" % self.data_dir, "a")
+        f.write("host replication %(username)s %(network)s md5" %
+                {"username": self.replication["username"], "network": self.replication["network"]})
+        f.close()
+
+
+    def write_recovery_conf(self, leader_hash):
+        leader = urlparse(leader_hash["address"])
+
+        f = open("%s/recovery.conf" % self.data_dir, "w")
+        f.write("""
+standby_mode = 'on'
+primary_slot_name = '%(recovery_slot)s'
+primary_conninfo = 'user=%(user)s password=%(password)s host=%(hostname)s port=%(port)s sslmode=prefer sslcompression=1'
+recovery_target_timeline = 'latest'
+""" % {"recovery_slot": self.name, "user": leader.username, "password": leader.password, "hostname": leader.hostname, "port": leader.port})
+        if "recovery_conf" in self.config:
+            for name, value in self.config["recovery_conf"].iteritems():
+                f.write("%s = '%s'" % (name, value))
+        f.close()
+
+    def follow_the_leader(self, leader_hash):
+        leader = urlparse(leader_hash["address"])
+        if os.system("grep 'host=%(hostname)s port=%(port)s' %(data_dir)s/recovery.conf" % {"hostname": leader.hostname, "port": leader.port, "data_dir": self.data_dir}) != 0:
+            self.write_recovery_conf(leader_hash);
+            self.restart()
+        return True
+
+    def promote(self):
+        return os.system("pg_ctl promote -w -D %s" % self.data_dir) == 0
+
+    def demote(self, leader):
+        self.write_recovery_conf(leader)
+        self.restart()
+
+    def create_replication_user(self):
+        self.query("CREATE USER \"%s\" WITH REPLICATION ENCRYPTED PASSWORD '%s';" % (self.replication["username"], self.replication["password"]))

postgres0.yml

@@ -0,0 +1,24 @@
+loop_wait: 10
+etcd:
+  scope: batman
+  ttl: 30
+  host: 127.0.0.1:4001
+postgresql:
+  name: postgresql0
+  listen: 127.0.0.1:5432
+  data_dir: data/postgresql0
+  replication:
+    username: replicator
+    password: rep-pass
+    network:  127.0.0.1/32
+  #recovery_conf:
+    #restore_command: cp ../wal_archive/%f %p
+  parameters:
+    archive_mode: "on"
+    wal_level: hot_standby
+    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
+    max_wal_senders: 5
+    wal_keep_segments: 8
+    archive_timeout: 1800s
+    max_replication_slots: 5
+    hot_standby: "on"

postgres1.yml

@@ -0,0 +1,24 @@
+loop_wait: 10
+etcd:
+  scope: batman
+  ttl: 30
+  host: 127.0.0.1:4001
+postgresql:
+  name: postgresql1
+  listen: 127.0.0.1:5433
+  data_dir: data/postgresql1
+  replication:
+    username: replicator
+    password: rep-pass
+    network: 127.0.0.1/32
+  #recovery_conf:
+    #restore_command: cp ../wal_archive/%f %p
+  parameters:
+    archive_mode: "on"
+    wal_level: hot_standby
+    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
+    max_wal_senders: 5
+    wal_keep_segments: 8
+    archive_timeout: 1800s
+    max_replication_slots: 5
+    hot_standby: "on"

run.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python
+
+import sys, os, yaml, time, urllib2, atexit
+from helpers.etcd import Etcd
+from helpers.postgresql import Postgresql
+from helpers.ha import Ha
+
+f = open(sys.argv[1], "r")
+config = yaml.load(f.read())
+f.close()
+
+print config
+
+etcd = Etcd(config["etcd"])
+postgresql = Postgresql(config["postgresql"])
+ha = Ha(postgresql, etcd)
+
+# stop postgresql on script exit
+def stop_postgresql():
+    postgresql.stop()
+atexit.register(stop_postgresql)
+
+# wait for etcd to be available
+etcd_ready = False
+while not etcd_ready:
+    try:
+        etcd.touch_member(postgresql.name, postgresql.connection_string)
+        etcd_ready = True
+    except urllib2.URLError:
+        print "waiting on etcd"
+        time.sleep(5)
+
+# is data directory empty?
+if postgresql.data_directory_empty():
+    # racing to initialize
+    if etcd.race("/initialize", postgresql.name):
+        postgresql.initialize()
+        etcd.take_leader(postgresql.name)
+        postgresql.start()
+        postgresql.create_replication_user()
+    else:
+        synced_from_leader = False
+        while not synced_from_leader:
+            leader = etcd.current_leader()
+            if leader == None:
+                time.sleep(5)
+                next
+            if postgresql.sync_from_leader(leader):
+                postgresql.write_recovery_conf(leader)
+                postgresql.start()
+                synced_from_leader = True
+            else:
+                time.sleep(5)
+else:
+    postgresql.write_recovery_conf({"address": "postgres://169.0.0.1:5432"})
+    postgresql.start()
+
+while True:
+    print ha.run_cycle()
+
+    # create replication slots
+    if postgresql.is_leader():
+        for node in etcd.get_client_path("/members?recursive=true")["node"]["nodes"]:
+            member = node["key"].split('/')[-1]
+            if member != postgresql.name:
+                postgresql.query("DO LANGUAGE plpgsql $$DECLARE somevar VARCHAR; BEGIN SELECT slot_name INTO somevar FROM pg_replication_slots WHERE slot_name = '%(slot)s' LIMIT 1; IF NOT FOUND THEN PERFORM pg_create_physical_replication_slot('%(slot)s'); END IF; END$$;" % {"slot": member})
+
+    time.sleep(config["loop_wait"])