Retry heartbeat after initial service registration (#247)

Sending a heartbeat to a not-yet registered service fails on both the
`discovery/consul` and `discovery/etcd` backends. Previously we'd catch
the error, register the service, and return so that we'd send a passing
heartbeat on the next health check poll. This change retries the
heartbeat on the newly registered service so that we become healthy more
quickly.

discovery/consul/consul.go

@@ -110,9 +110,16 @@ func (c *Consul) SendHeartbeat(service *discovery.ServiceDefinition) {
 		log.Infof("%v\nService not registered, registering...", err)
 		if err = c.registerService(*service); err != nil {
 			log.Warnf("Service registration failed: %s", err)
+			return
 		}
 		if err = c.registerCheck(*service); err != nil {
 			log.Warnf("Check registration failed: %s", err)
+			return
+		}
+		// now that we're ensured we're registered, we can push the
+		// heartbeat again
+		if err := c.Agent().PassTTL(service.ID, "ok"); err != nil {
+			log.Errorf("Failed to write heartbeat: %s", err)
 		}
 	}
 }

discovery/consul/consul_test.go

@@ -59,16 +59,9 @@ func TestConsulTTLPass(t *testing.T) {
 	consul, service := setupConsul("service-TestConsulTTLPass")
 	id := service.ID
 
-	consul.SendHeartbeat(service) // force registration
+	consul.SendHeartbeat(service) // force registration and 1st heartbeat
 	checks, _ := consul.Agent().Checks()
 	check := checks[id]
-	if check.Status != "critical" {
-		t.Fatalf("status of check %s should be 'critical' but is %s", id, check.Status)
-	}
-
-	consul.SendHeartbeat(service) // write TTL and verify
-	checks, _ = consul.Agent().Checks()
-	check = checks[id]
 	if check.Status != "passing" {
 		t.Fatalf("status of check %s should be 'passing' but is %s", id, check.Status)
 	}

discovery/etcd/etcd.go

@@ -103,6 +103,12 @@ func (c *Etcd) SendHeartbeat(service *discovery.ServiceDefinition) {
 		log.Infof("Service not registered, registering...")
 		if err := c.registerService(service); err != nil {
 			log.Warnf("Error registering service %s: %s", service.Name, err)
+			return
+		}
+		// now that we're ensured we're registered, we can push the
+		// heartbeat again
+		if err := c.updateServiceTTL(service); err != nil {
+			log.Errorf("Failed to write heartbeat: %s", err)
 		}
 	}
 }

discovery/etcd/etcd_test.go

@@ -40,22 +40,16 @@ func TestEtcdParseStringEndpoints(t *testing.T) {
 	}
 }
 
-func TestEtcdTTLPass(t *testing.T) {
+func TestEtcdTTLExpires(t *testing.T) {
 	etcd, service := setupEtcd("service-TestEtcdTTLPass")
 	id := service.ID
 
-	etcd.SendHeartbeat(service) // force registration
-	if !checkServiceExists(etcd, service) {
-		t.Fatalf("Expected service %s to be registered, but was not", id)
-	}
-
-	etcd.SendHeartbeat(service) // write TTL and verify
-	if !checkServiceExists(etcd, service) {
-		t.Fatalf("Expected service %s to be registered, but was not", id)
+	etcd.SendHeartbeat(service) // force registration and TTL
+	if !checkServiceHealthy(etcd, service) {
+		t.Fatalf("Expected service %s to be registered and healthy, but was not", id)
 	}
 
-	time.Sleep(2 * time.Second)
-
+	time.Sleep(2 * time.Second) // wait for TTL to expire
 	if checkServiceExists(etcd, service) {
 		t.Fatalf("Expected service %s to be deregistered registered", id)
 	}
@@ -90,8 +84,7 @@ func TestEtcdCheckForChanges(t *testing.T) {
 	if etcd.CheckForUpstreamChanges(backend, "") {
 		t.Fatalf("First read of %s should show `false` for change", id)
 	}
-	etcd.SendHeartbeat(service) // force registration
-	etcd.SendHeartbeat(service) // write TTL
+	etcd.SendHeartbeat(service) // force registration and TTL
 
 	if !etcd.CheckForUpstreamChanges(backend, "") {
 		t.Errorf("%v should have changed after first health check TTL", id)
@@ -114,3 +107,17 @@ func checkServiceExists(etcd *Etcd, service *discovery.ServiceDefinition) bool {
 	}
 	return true
 }
+
+func checkServiceHealthy(etcd *Etcd, service *discovery.ServiceDefinition) bool {
+	key := etcd.getNodeKey(service)
+	if resp, err := etcd.API.Get(context.Background(), key, nil); err != nil {
+		if etcdErr, ok := err.(client.Error); ok {
+			return etcdErr.Code != client.ErrorCodeKeyNotFound
+		}
+	} else {
+		if len(resp.Node.Nodes) == 1 {
+			return true
+		}
+	}
+	return false
+}