Initial commit

lease/heart.go

@@ -0,0 +1,34 @@
+package lease
+
+import "time"
+
+// Heart contains a lease that survives when heartbeats happen within a given ttl interval
+type Heart struct {
+	lastBeat time.Time
+	ttl      time.Duration
+}
+
+// NewHeart creates a new heart instance
+func NewHeart(ttl int64) *Heart {
+	ttlDuration := time.Nanosecond * time.Duration(ttl)
+	return &Heart{
+		lastBeat: time.Now(),
+		ttl:      ttlDuration,
+	}
+}
+
+// Valid tries to beat the heart, true if still alive, false if dead
+func (h *Heart) Valid() bool {
+	if !h.Check() {
+		return false
+	}
+
+	// update the last beat to now
+	h.lastBeat = time.Now()
+	return true
+}
+
+// Check if the heartbeat is valid
+func (h *Heart) Check() bool {
+	return time.Since(h.lastBeat) <= h.ttl
+}

lease/heart_test.go

@@ -0,0 +1,52 @@
+package lease
+
+import (
+	"testing"
+	"time"
+)
+
+func TestHeartCheck(t *testing.T) {
+	heart := NewHeart(int64(time.Second))
+
+	if !heart.Check() {
+		t.Error("Heartbeat check failed first check too soon")
+	}
+
+	time.Sleep(time.Second + time.Nanosecond)
+
+	if heart.Check() {
+		t.Error("Heartbeat check succeeded after its ttl")
+	}
+}
+
+func TestHeartValid(t *testing.T) {
+	heart := NewHeart(int64(time.Second))
+
+	if !heart.Valid() {
+		t.Error("Heartbeat valid failed first check too soon")
+	}
+
+	time.Sleep(time.Second / 3)
+
+	if !heart.Valid() {
+		t.Error("Heartbeat valid failed second check too soon")
+	}
+
+	time.Sleep(time.Second / 2)
+
+	if !heart.Valid() {
+		t.Error("Hearbeat valid failed third check too soon")
+	}
+
+	time.Sleep(time.Second / 2)
+
+	if !heart.Valid() {
+		t.Error("Heartbeat valid failed fourth check too soon")
+	}
+
+	time.Sleep(time.Second)
+
+	if heart.Valid() {
+		t.Error("Heartbeat valid succeeded after its ttl")
+	}
+}

lease/leaser.go

@@ -0,0 +1,10 @@
+package lease
+
+// Leaser denotes whether a given lease is still valid
+type Leaser interface {
+	// Update and check the leaser
+	Valid() bool
+
+	// Check the leaser without any sort of updates
+	Check() bool
+}

lease/manager.go

@@ -0,0 +1,94 @@
+package lease
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/tblyler/sheepmq/shepard"
+)
+
+var (
+	// ErrLeased denotes the item is already leased
+	ErrLeased = errors.New("Item is already leased")
+
+	// ErrNoLeaser denotes no leaser was provided to add a lease
+	ErrNoLeaser = errors.New("No leaser provided")
+)
+
+// Manager contains many leases and their validity
+type Manager struct {
+	leases map[uint64]Leaser
+	locker sync.RWMutex
+}
+
+// NewManager creates a new Manager instance
+func NewManager() *Manager {
+	return &Manager{
+		leases: make(map[uint64]Leaser),
+	}
+}
+
+// AddLease to the manager for the given item
+func (m *Manager) AddLease(id uint64, info *shepard.GetInfo) error {
+	if m.CheckLease(id) {
+		return ErrLeased
+	}
+
+	var leaser Leaser
+	if info.TimeoutLease != nil {
+		leaser = NewTimeout(info.TimeoutLease.Ttl)
+	} else if info.PidLease != nil {
+		leaser = NewPID(int(info.PidLease.Pid))
+	} else if info.HeartLease != nil {
+		leaser = NewHeart(info.HeartLease.Ttl)
+	} else {
+		return ErrNoLeaser
+	}
+
+	m.locker.Lock()
+	m.leases[id] = leaser
+	m.locker.Unlock()
+
+	return nil
+}
+
+// CheckLease for validity
+func (m *Manager) CheckLease(id uint64) bool {
+	m.locker.RLock()
+	lease, exists := m.leases[id]
+	if !exists {
+		m.locker.RUnlock()
+		return false
+	}
+
+	ret := lease.Valid()
+	m.locker.RUnlock()
+	if !ret {
+		// delete the lease since it is no longer valid
+		m.locker.Lock()
+		delete(m.leases, id)
+		m.locker.Unlock()
+	}
+
+	return ret
+}
+
+// PruneLeases that fail their checks
+func (m *Manager) PruneLeases() {
+	deleteKeys := []uint64{}
+	m.locker.RLock()
+
+	for key, lease := range m.leases {
+		if !lease.Check() {
+			deleteKeys = append(deleteKeys, key)
+		}
+	}
+
+	m.locker.RUnlock()
+	m.locker.Lock()
+	for _, key := range deleteKeys {
+		delete(m.leases, key)
+	}
+
+	m.locker.Unlock()
+}

lease/manager_test.go

@@ -0,0 +1,134 @@
+package lease
+
+import (
+	"testing"
+	"time"
+
+	"github.com/tblyler/sheepmq/shepard"
+)
+
+func TestManagerAddCheckLease(t *testing.T) {
+	manager := NewManager()
+
+	info := &shepard.GetInfo{}
+
+	err := manager.AddLease(123, info)
+	if err != ErrNoLeaser {
+		t.Error("Failed to get ErrNoLeaser on add lease got", err)
+	}
+
+	info.PidLease = &shepard.PidLease{
+		// use a bad PID
+		Pid: 0,
+	}
+
+	err = manager.AddLease(123, info)
+	if err != nil {
+		t.Error("Failed to add a crappy PID leaser", err)
+	}
+
+	err = manager.AddLease(123, info)
+	if err != nil {
+		t.Error("Failed to add a crappy PID leaser ontop of another", err)
+	}
+
+	info.PidLease = nil
+
+	info.HeartLease = &shepard.HeartbeatLease{
+		Ttl: int64(time.Second / 2),
+	}
+
+	err = manager.AddLease(2, info)
+	if err != nil {
+		t.Error("Failed to add a valid heartbeat lease", err)
+	}
+
+	if !manager.CheckLease(2) {
+		t.Error("Failed to check for valid heartbeat lease")
+	}
+
+	info.HeartLease = nil
+	info.TimeoutLease = &shepard.TimeLease{
+		Ttl: int64(time.Second),
+	}
+
+	err = manager.AddLease(123, info)
+	if err != nil {
+		t.Error("Failed to add a good timeout lease of 1 second")
+	}
+
+	err = manager.AddLease(123, info)
+	if err == nil {
+		t.Error("Should not be able to add a lease ontop of another valid one")
+	}
+
+	err = manager.AddLease(124, info)
+	if err != nil {
+		t.Error("Failed to add a valid lease against a different id", err)
+	}
+
+	if !manager.CheckLease(123) {
+		t.Error("Failed to make sure valid lease was valid")
+	}
+
+	if !manager.CheckLease(124) {
+		t.Error("Failed to make sure valid lease was valid")
+	}
+
+	time.Sleep(time.Second)
+
+	if manager.CheckLease(123) {
+		t.Error("failed to make sure invalid lease was invalid")
+	}
+
+	if manager.CheckLease(124) {
+		t.Error("failed to make sure invalid lease was invalid")
+	}
+}
+
+func TestManagerPruneLeases(t *testing.T) {
+	manager := NewManager()
+
+	info := &shepard.GetInfo{}
+	info.PidLease = &shepard.PidLease{
+		Pid: 0,
+	}
+
+	err := manager.AddLease(5, info)
+	if err != nil {
+		t.Error("Failed to add crappy PID lease", err)
+	}
+
+	info.PidLease = nil
+	info.TimeoutLease = &shepard.TimeLease{
+		Ttl: int64(time.Second),
+	}
+
+	err = manager.AddLease(2, info)
+	if err != nil {
+		t.Error("Failed to add valid timeout lease", err)
+	}
+
+	info.TimeoutLease = nil
+	info.HeartLease = &shepard.HeartbeatLease{
+		Ttl: int64(time.Second),
+	}
+
+	err = manager.AddLease(1, info)
+	if err != nil {
+		t.Error("Failed to add valid heartbeat lease", err)
+	}
+
+	manager.PruneLeases()
+
+	if len(manager.leases) != 2 {
+		t.Errorf("Should have 2 leases left, have %d", len(manager.leases))
+	}
+
+	time.Sleep(time.Second)
+
+	manager.PruneLeases()
+	if len(manager.leases) != 0 {
+		t.Errorf("Should have 0 leases left, have %d", len(manager.leases))
+	}
+}

lease/pid.go

@@ -0,0 +1,34 @@
+package lease
+
+import (
+	"os"
+	"syscall"
+)
+
+// PID contains a lease that is valid until the given PID no longer exists
+type PID struct {
+	pid int
+}
+
+// NewPID creates a new PID leaser instance
+func NewPID(pid int) *PID {
+	return &PID{
+		pid: pid,
+	}
+}
+
+// Valid checks if the PID still exists
+func (p *PID) Valid() bool {
+	return p.Check()
+}
+
+// Check if the PID still exists
+func (p *PID) Check() bool {
+	process, err := os.FindProcess(p.pid)
+	if err != nil {
+		return false
+	}
+
+	// if nil, PID exists
+	return process.Signal(syscall.Signal(0)) == nil
+}

lease/pid_test.go

@@ -0,0 +1,31 @@
+package lease
+
+import (
+	"os/exec"
+	"testing"
+)
+
+func TestPIDValid(t *testing.T) {
+	// create a sleep command for 1 second
+	cmd := exec.Command("sleep", "1")
+	err := cmd.Start()
+	if err != nil {
+		t.Fatal("Failed to create a sleep process:", err)
+	}
+
+	pid := NewPID(cmd.Process.Pid)
+	if !pid.Valid() {
+		t.Error("PID died too soon")
+	}
+
+	cmd.Wait()
+
+	if pid.Valid() {
+		t.Error("PID didn't die somehow")
+	}
+
+	pid = NewPID(-1)
+	if pid.Valid() {
+		t.Error("Negative PIDS are not a thing")
+	}
+}

lease/timeout.go

@@ -0,0 +1,27 @@
+package lease
+
+import (
+	"time"
+)
+
+// Timeout expires after a given timeout
+type Timeout struct {
+	eol time.Time
+}
+
+// NewTimeout creates a new timeout instance
+func NewTimeout(ttl int64) *Timeout {
+	return &Timeout{
+		eol: time.Now().Add(time.Nanosecond * time.Duration(ttl)),
+	}
+}
+
+// Valid Determines whether the timeout has been reached
+func (t *Timeout) Valid() bool {
+	return t.Check()
+}
+
+// Check if the timeout has been reached
+func (t *Timeout) Check() bool {
+	return time.Now().Before(t.eol)
+}

lease/timeout_test.go

@@ -0,0 +1,26 @@
+package lease
+
+import (
+	"testing"
+	"time"
+)
+
+func TestTimeoutValid(t *testing.T) {
+	timeout := NewTimeout(int64(time.Second))
+
+	if !timeout.Valid() {
+		t.Error("timeout valid failed too soon")
+	}
+
+	time.Sleep(time.Nanosecond * 100)
+
+	if !timeout.Valid() {
+		t.Error("timeout valid failed too soon")
+	}
+
+	time.Sleep(time.Second - (time.Nanosecond * 100))
+
+	if timeout.Valid() {
+		t.Error("timeout valid should not be succeeding")
+	}
+}