From 854a4623c350ca2e71741cc93142f13e2b6afe8a Mon Sep 17 00:00:00 2001
From: Xinzhao Xu <xuxinzhao@caicloud.io>
Date: Fri, 23 Aug 2019 14:13:50 +0800
Subject: [PATCH] feat(plugin): add a type parameter for the health check
 plugin (#286)

* feat(plugin): add a type parameter for the health check plugin

* feat(plugin): HealthCheckerWithType

* feat(plugin): fix ci

* feat(plugin): remove NewCheckerWithType
---
 examples/getting-started/healthcheck/echo.go | 89 ++++++++++++++++++++
 plugins/healthcheck/healthz.go               | 49 +++++++++--
 2 files changed, 133 insertions(+), 5 deletions(-)
 create mode 100644 examples/getting-started/healthcheck/echo.go

diff --git a/examples/getting-started/healthcheck/echo.go b/examples/getting-started/healthcheck/echo.go
new file mode 100644
index 00000000..96e5b826
--- /dev/null
+++ b/examples/getting-started/healthcheck/echo.go
@@ -0,0 +1,89 @@
+/*
+Copyright 2018 Caicloud Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+import (
+	"context"
+
+	"github.com/caicloud/nirvana"
+	"github.com/caicloud/nirvana/config"
+	"github.com/caicloud/nirvana/definition"
+	"github.com/caicloud/nirvana/errors"
+	"github.com/caicloud/nirvana/log"
+	"github.com/caicloud/nirvana/plugins/healthcheck"
+)
+
+var echo = definition.Descriptor{
+	Path:        "/echo",
+	Description: "Echo API",
+	Definitions: []definition.Definition{
+		{
+			Method:   definition.Get,
+			Function: Echo,
+			Consumes: []string{definition.MIMEAll},
+			Produces: []string{definition.MIMEJSON},
+			Parameters: []definition.Parameter{
+				{
+					Source:      definition.Query,
+					Name:        "msg",
+					Description: "Corresponding to the second parameter",
+				},
+			},
+			Results: []definition.Result{
+				{
+					Destination: definition.Data,
+					Description: "Corresponding to the first result",
+				},
+				{
+					Destination: definition.Error,
+					Description: "Corresponding to the second result",
+				},
+			},
+		},
+	},
+}
+
+// API function.
+func Echo(ctx context.Context, msg string) (string, error) {
+	return msg, nil
+}
+
+func main() {
+	cmd := config.NewDefaultNirvanaCommand()
+	cfg := nirvana.NewDefaultConfig()
+	cfg.Configure(
+		nirvana.Descriptor(echo),
+		healthcheck.CheckerWithType(func(ctx context.Context, checkType string) error {
+			switch checkType {
+			case healthcheck.LivenessCheck:
+				// do something
+				return nil
+			case healthcheck.ReadinessCheck:
+				// do something
+				return nil
+			default:
+				return errors.BadRequest.Build("error", "unknown type ${type}").Error(checkType)
+			}
+		}),
+		// healthcheck.Checker(func(ctx context.Context) error {
+		// 	return nil
+		// }),
+	)
+	if err := cmd.ExecuteWithConfig(cfg); err != nil {
+		log.Fatal(err)
+	}
+}
diff --git a/plugins/healthcheck/healthz.go b/plugins/healthcheck/healthz.go
index c765ab63..18fedfe1 100644
--- a/plugins/healthcheck/healthz.go
+++ b/plugins/healthcheck/healthz.go
@@ -35,13 +35,29 @@ func defaultHealthChecker(ctx context.Context) error {
 	return nil
 }
 
+// HealthCheckerWithType checks if current server is healthy.
+// The `checkType` parameter indicates the type of health check, such as liveness or readiness.
+type HealthCheckerWithType func(ctx context.Context, checkType string) error
+
+func defaultHealthCheckerWithType(ctx context.Context, checkType string) error {
+	return nil
+}
+
+const (
+	// LivenessCheck represents liveness check.
+	LivenessCheck = "liveness"
+	// ReadinessCheck represents readiness check.
+	ReadinessCheck = "readiness"
+)
+
 // ExternalConfigName is the external config name of health check.
 const ExternalConfigName = "healthcheck"
 
 // config is healthcheck config.
 type config struct {
-	path    string
-	checker HealthChecker
+	path            string
+	checker         HealthChecker
+	checkerWithType HealthCheckerWithType
 }
 
 type healthcheckInstaller struct{}
@@ -55,14 +71,24 @@ func (i *healthcheckInstaller) Name() string {
 func (i *healthcheckInstaller) Install(builder service.Builder, cfg *nirvana.Config) error {
 	var err error
 	wrapper(cfg, func(c *config) {
+		var parameters []definition.Parameter
+		var function interface{} = c.checker
+		if c.checkerWithType != nil {
+			parameters = []definition.Parameter{
+				definition.QueryParameterFor("type", "the type of health check"),
+			}
+			function = c.checkerWithType
+		}
+
 		err = builder.AddDescriptor(definition.Descriptor{
 			Path:     c.path,
 			Consumes: []string{definition.MIMEAll},
 			Produces: []string{definition.MIMEAll},
 			Definitions: []definition.Definition{{
-				Method:   definition.Get,
-				Results:  []definition.Result{definition.ErrorResult()},
-				Function: c.checker,
+				Method:     definition.Get,
+				Results:    []definition.Result{definition.ErrorResult()},
+				Parameters: parameters,
+				Function:   function,
 			}},
 		})
 	})
@@ -108,6 +134,19 @@ func Checker(checker HealthChecker) nirvana.Configurer {
 	}
 }
 
+// CheckerWithType returns a configurer to set health checker with a type parameter.
+func CheckerWithType(checker HealthCheckerWithType) nirvana.Configurer {
+	if checker == nil {
+		checker = defaultHealthCheckerWithType
+	}
+	return func(c *nirvana.Config) error {
+		wrapper(c, func(c *config) {
+			c.checkerWithType = checker
+		})
+		return nil
+	}
+}
+
 func wrapper(c *nirvana.Config, f func(c *config)) {
 	conf := c.Config(ExternalConfigName)
 	var cfg *config