diff --git a/consul/state/state_store.go b/consul/state/state_store.go index 636b9a21d..085a1c8f0 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 db46d42e1..9137bd6ce 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) + } +}