From cf569f48486aac10dad456973a5b43d7100fccbb Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Mon, 24 Aug 2015 20:51:07 -0700
Subject: [PATCH] consul/state: persisting health checks works

---
 consul/state/state_store.go      | 73 ++++++++++++++++++++++++++++++++
 consul/state/state_store_test.go | 40 +++++++++++++++++
 2 files changed, 113 insertions(+)

diff --git a/consul/state/state_store.go b/consul/state/state_store.go
index 636b9a21d7..085a1c8f0e 100644
--- a/consul/state/state_store.go
+++ b/consul/state/state_store.go
@@ -1,6 +1,7 @@
 package state
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -9,6 +10,16 @@ import (
 	"github.com/hashicorp/go-memdb"
 )
 
+var (
+	// ErrMissingNode is the error returned when trying an operation
+	// which requires a node registration but none exists.
+	ErrMissingNode = errors.New("Missing node registration")
+
+	// ErrMissingService is the error we return if trying an
+	// operation which requires a service but none exists.
+	ErrMissingService = errors.New("Missing service registration")
+)
+
 // StateStore is where we store all of Consul's state, including
 // records of node registrations, services, checks, key/value
 // pairs and more. The DB is entirely in-memory and is constructed
@@ -322,3 +333,65 @@ func (s *StateStore) deleteNodeServiceTxn(idx uint64, nodeID, serviceID string,
 	// TODO: watch trigger
 	return nil
 }
+
+// EnsureCheck is used to store a check registration in the db.
+func (s *StateStore) EnsureCheck(idx uint64, hc *structs.HealthCheck) error {
+	tx := s.db.Txn(true)
+	defer tx.Abort()
+
+	// Call the check registration
+	if err := s.ensureCheckTxn(idx, hc, tx); err != nil {
+		return err
+	}
+
+	tx.Commit()
+	return nil
+}
+
+// ensureCheckTransaction is used as the inner method to handle inserting
+// a health check into the state store. It ensures safety against inserting
+// checks with no matching node or service.
+func (s *StateStore) ensureCheckTxn(idx uint64, hc *structs.HealthCheck, tx *memdb.Txn) error {
+	// Use the default check status if none was provided
+	if hc.Status == "" {
+		hc.Status = structs.HealthCritical
+	}
+
+	// Get the node
+	node, err := tx.First("nodes", "id", hc.Node)
+	if err != nil {
+		return fmt.Errorf("failed node lookup: %s", err)
+	}
+	if node == nil {
+		return ErrMissingNode
+	}
+
+	// If the check is associated with a service, check that we have
+	// a registration for the service.
+	if hc.ServiceID != "" {
+		service, err := tx.First("services", "id", hc.Node, hc.ServiceID)
+		if err != nil {
+			return fmt.Errorf("failed service lookup: %s", err)
+		}
+		if service == nil {
+			return ErrMissingService
+		}
+
+		// Copy in the service name
+		hc.ServiceName = service.(*structs.ServiceNode).ServiceName
+	}
+
+	// TODO: invalidate sessions if status == critical
+
+	// Persist the check registration in the db
+	if err := tx.Insert("services", hc); err != nil {
+		return fmt.Errorf("failed inserting service: %s", err)
+	}
+	if err := tx.Insert("index", &IndexEntry{"checks", idx}); err != nil {
+		return fmt.Errorf("failed updating index: %s", err)
+	}
+
+	// TODO: trigger watches
+
+	return nil
+}
diff --git a/consul/state/state_store_test.go b/consul/state/state_store_test.go
index db46d42e19..9137bd6ce8 100644
--- a/consul/state/state_store_test.go
+++ b/consul/state/state_store_test.go
@@ -274,3 +274,43 @@ func TestStateStore_DeleteNodeService(t *testing.T) {
 		t.Fatalf("bad index: %d", idx)
 	}
 }
+
+func TestStateStore_EnsureCheck(t *testing.T) {
+	s := testStateStore(t)
+
+	// Create a node and insert it
+	node := &structs.Node{
+		Node:    "node1",
+		Address: "1.1.1.1",
+	}
+	if err := s.EnsureNode(1, node); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Create a service and insert it
+	service := &structs.NodeService{
+		ID:      "service1",
+		Service: "redis",
+		Tags:    []string{"prod"},
+		Address: "1.1.1.1",
+		Port:    1111,
+	}
+	if err := s.EnsureService(2, "node1", service); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Create a check associated with the node and insert it
+	check := &structs.HealthCheck{
+		Node:        "node1",
+		CheckID:     "check1",
+		Name:        "redis check",
+		Status:      structs.HealthPassing,
+		Notes:       "test check",
+		Output:      "aaa",
+		ServiceID:   "service1",
+		ServiceName: "redis",
+	}
+	if err := s.EnsureCheck(3, check); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+}