Merge pull request #55 from lynes-io/readiness-probe

Add readiness probe support for consul check

README.md

@@ -91,6 +91,9 @@ There are available annotations which can be used as pod's annotations.
 |`consul.register/service.name`|`service_name`|Determine name of service in Consul. If not given then is used the name of resource which created the POD. Only available if `register_source` is set on `pod`|
 |`consul.register/service.meta.<key>`|`<value>`|Adds `key`/`value` service meta. Eg. `"consul.register/service.meta.redis_version"`=`"4.0"` results in meta `redis_version=4.0`|
 |`consul.register/pod.container.name`|`container_name`|Container name or list of names (next name should be separated by comma) which will be taken into account. If omitted, all containers in POD will be registered|
+|`consul.register/pod.container.probe.liveness`|`true`\|`false`|Use container `Liveness probe` for checks. Default is `true`.
+|`consul.register/pod.container.probe.readiness`|`true`\|`false`|Use container `Readiness probe` for checks. Default is `false`|
+
 
 The example of how to use annotation you can see [here](https://github.com/tczekajlo/kube-consul-register/blob/master/examples/nginx.yaml).
 

controller/pods/controller.go

@@ -38,6 +38,8 @@ const (
 	ConsulRegisterServiceMetaPrefixAnnotation string = "consul.register/service.meta."
 	CreatedByAnnotation                       string = "kubernetes.io/created-by"
 	ExpectedContainerNamesAnnotation          string = "consul.register/pod.container.name"
+  ContainerProbeLivenessAnnotation          string = "consul.register/pod.container.probe.liveness"
+	ContainerProbeReadinessAnnotation         string = "consul.register/pod.container.probe.readiness"
 )
 
 var (
@@ -446,9 +448,15 @@ func (p *PodInfo) PodToConsulService(containerStatus v1.ContainerStatus, cfg *co
 		return service, fmt.Errorf("Port's equal to 0")
 	}
 	service.Port = port
-	service.Check = p.livenessProbeToConsulCheck(p.getContainerLivenessProbe(containerStatus.Name))
 	service.Address = p.IP
 
+	if p.isProbeLivenessEnabled() {
+		service.Checks = append(service.Checks, p.probeToConsulCheck(p.getContainerLivenessProbe(containerStatus.Name), "Liveness Probe"))
+	}
+	if p.isProbeReadinessEnabled() {
+		service.Checks = append(service.Checks, p.probeToConsulCheck(p.getContainerReadinessProbe(containerStatus.Name), "Readiness Probe"))
+	}
+
 	return service, nil
 }
 
@@ -471,6 +479,44 @@ func (p *PodInfo) isRegisterEnabled() bool {
 	return true
 }
 
+func (p *PodInfo) isProbeLivenessEnabled() bool {
+	// Default if not set should be true
+	if value, ok := p.Annotations[ContainerProbeLivenessAnnotation]; ok {
+		enabled, err := strconv.ParseBool(value)
+		if err != nil {
+			glog.Errorf("Can't convert value of %s annotation: %s", ContainerProbeLivenessAnnotation, err)
+			return true
+		}
+		if enabled {
+			glog.Infof("Pod %s in %s namespace has liveness probe enabled by annotation. Value: %s", p.Name, p.Namespace, value)
+			return true
+		} else {
+			glog.Infof("Pod %s in %s namespace has liveness probe disabled by annotation. Value: %s", p.Name, p.Namespace, value)
+			return false
+		}
+	}
+	return true
+}
+
+func (p *PodInfo) isProbeReadinessEnabled() bool {
+	// Default if not set should be false
+	if value, ok := p.Annotations[ContainerProbeReadinessAnnotation]; ok {
+		enabled, err := strconv.ParseBool(value)
+		if err != nil {
+			glog.Errorf("Can't convert value of %s annotation: %s", ContainerProbeReadinessAnnotation, err)
+			return false
+		}
+		if enabled {
+			glog.Infof("Pod %s in %s namespace has readiness probe enabled by annotation. Value: %s", p.Name, p.Namespace, value)
+			return true
+		} else {
+			glog.Infof("Pod %s in %s namespace has readiness probe disabled by annotation. Value: %s", p.Name, p.Namespace, value)
+			return false
+		}
+	}
+	return false
+}
+
 func (p *PodInfo) expectedContainerNames(containerName string) bool {
 	if value, ok := p.Annotations[ExpectedContainerNamesAnnotation]; ok {
 		for _, name := range strings.Split(value, ",") {
@@ -485,7 +531,7 @@ func (p *PodInfo) expectedContainerNames(containerName string) bool {
 	return false
 }
 
-func (p *PodInfo) livenessProbeToConsulCheck(probe *v1.Probe) *consulapi.AgentServiceCheck {
+func (p *PodInfo) probeToConsulCheck(probe *v1.Probe, probeName string) *consulapi.AgentServiceCheck {
 	check := &consulapi.AgentServiceCheck{}
 
 	if probe == nil {
@@ -496,6 +542,7 @@ func (p *PodInfo) livenessProbeToConsulCheck(probe *v1.Probe) *consulapi.AgentSe
 		return check
 	}
 
+	check.Name = probeName
 	check.Status = "passing"
 	check.Interval = fmt.Sprintf("%ds", probe.PeriodSeconds)
 	check.Timeout = fmt.Sprintf("%ds", probe.TimeoutSeconds)
@@ -523,6 +570,15 @@ func (p *PodInfo) getContainerLivenessProbe(searchContainer string) *v1.Probe {
 	return nil
 }
 
+func (p *PodInfo) getContainerReadinessProbe(searchContainer string) *v1.Probe {
+	for _, container := range p.Containers {
+		if container.Name == searchContainer {
+			return container.ReadinessProbe
+		}
+	}
+	return nil
+}
+
 func (p *PodInfo) labelsToTags(containerName string) []string {
 	var tags []string
 	tags = append(tags, p.Name)

controller/pods/controller_test.go

@@ -70,7 +70,7 @@ func TestPodInfoMethods(t *testing.T) {
 	assert.Equal(t, true, isEnabledByAnnotation)
 }
 
-func TestLivenessProbeToConsulCheck(t *testing.T) {
+func TestProbeToConsulCheck(t *testing.T) {
 	t.Parallel()
 	emptyCheck := consulapi.AgentServiceCheck{}
 
@@ -102,11 +102,12 @@ func TestLivenessProbeToConsulCheck(t *testing.T) {
 		},
 	}
 
-	httpCheck := podInfo.livenessProbeToConsulCheck(httpProbe)
-	tcpCheck := podInfo.livenessProbeToConsulCheck(tcpProbe)
-	noProbeCheck := podInfo.livenessProbeToConsulCheck(nil)
-	execCheck := podInfo.livenessProbeToConsulCheck(execProbe)
+	httpCheck := podInfo.probeToConsulCheck(httpProbe, "Liveness Probe")
+	tcpCheck := podInfo.probeToConsulCheck(tcpProbe, "Liveness Probe")
+	noProbeCheck := podInfo.probeToConsulCheck(nil, "Liveness Probe")
+	execCheck := podInfo.probeToConsulCheck(execProbe, "Liveness Probe")
 
+	assert.Equal(t, "Liveness Probe", httpCheck.Name)
 	assert.Equal(t, "http://192.168.8.8:8080/ping", httpCheck.HTTP)
 	assert.Equal(t, "192.168.8.8:5432", tcpCheck.TCP)
 	assert.Equal(t, emptyCheck, *noProbeCheck)