From a957e90ca867d78376a2a156b54e726b7c65c6eb Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Thu, 5 Dec 2019 17:15:35 +0300
Subject: [PATCH] add RegisterCheck web server option for internal health
 checks

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 web/options.go | 14 ++++++++++++++
 web/service.go |  7 +++++++
 web/web.go     |  4 +++-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/web/options.go b/web/options.go
index 9a31f6d0..6f1233db 100644
--- a/web/options.go
+++ b/web/options.go
@@ -25,6 +25,9 @@ type Options struct {
 	RegisterTTL      time.Duration
 	RegisterInterval time.Duration
 
+	// RegisterCheck runs a check function before registering the service
+	RegisterCheck func(context.Context) error
+
 	Server  *http.Server
 	Handler http.Handler
 
@@ -62,6 +65,10 @@ func newOptions(opts ...Option) Options {
 		o(&opt)
 	}
 
+	if opt.RegisterCheck == nil {
+		opt.RegisterCheck = DefaultRegisterCheck
+	}
+
 	return opt
 }
 
@@ -225,3 +232,10 @@ func StaticDir(d string) Option {
 		o.StaticDir = d
 	}
 }
+
+// RegisterCheck run func before registry service
+func RegisterCheck(fn func(context.Context) error) Option {
+	return func(o *Options) {
+		o.RegisterCheck = fn
+	}
+}
diff --git a/web/service.go b/web/service.go
index 49850d9e..e9d6801b 100644
--- a/web/service.go
+++ b/web/service.go
@@ -123,6 +123,13 @@ func (s *service) register() error {
 	srv := s.genSrv()
 	srv.Endpoints = s.srv.Endpoints
 	s.srv = srv
+
+	// use RegisterCheck func before register
+	if err := s.opts.RegisterCheck(s.opts.Context); err != nil {
+		log.Logf("Server %s-%s register check error: %s", s.opts.Name, s.opts.Id, err)
+		return err
+	}
+
 	return r.Register(s.srv, registry.RegisterTTL(s.opts.RegisterTTL))
 }
 
diff --git a/web/web.go b/web/web.go
index f88e25ef..5c06c2ae 100644
--- a/web/web.go
+++ b/web/web.go
@@ -2,6 +2,7 @@
 package web
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -32,7 +33,8 @@ var (
 	DefaultRegisterInterval = time.Second * 30
 
 	// static directory
-	DefaultStaticDir = "html"
+	DefaultStaticDir     = "html"
+	DefaultRegisterCheck = func(context.Context) error { return nil }
 )
 
 // NewService returns a new web.Service