Repository: go-kratos/kratos
Pull request #3131 (2023-12-19 07:44:56)

feat: add server option for NotFoundHandler and MethodNotAllowedHandler
#### Description (what this PR does / why we need it):
pprof will be registered to DefaultServeMux by default, and go-kratos's NotFoundHandler and MethodNotAllowedHandler will use DefaultServeMux, resulting in the inability to flexibly control the switch of pprof. 
This PR adds ServerOption to modify NotFoundHandler and MethodNotAllowedHandler

#### Which issue(s) this PR fixes (resolves / be part of):

#### Other special notes for the reviewers:



codecov-commenter: ## [Codecov](https://app.codecov.io/gh/go-kratos/kratos/pull/3131?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`856bc9a`)](https://app.codecov.io/gh/go-kratos/kratos/commit/856bc9a17b4a62dc7f797bf2b1449ed053cc9d50?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) 84.43% compared to head [(`85abb47`)](https://app.codecov.io/gh/go-kratos/kratos/pull/3131?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) 84.55%.
> Report is 1 commits behind head on main.

> :exclamation: Current head 85abb47 differs from pull request most recent head 36b648d. Consider uploading reports for the commit 36b648d to get more accurate results


:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #3131      +/-   ##
==========================================
+ Coverage   84.43%   84.55%   +0.12%     
==========================================
  Files          88       88              
  Lines        3983     3989       +6     
==========================================
+ Hits         3363     3373      +10     
+ Misses        443      441       -2     
+ Partials      177      175       -2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/go-kratos/kratos/pull/3131?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).


Patch:
diff --git a/contrib/registry/consul/client.go b/contrib/registry/consul/client.go
index ea5fe846..ab1551cd 100644
--- a/contrib/registry/consul/client.go
+++ b/contrib/registry/consul/client.go
@@ -220,7 +220,11 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 						_ = c.cli.Agent().ServiceDeregister(svc.ID)
 						return
 					}
-					err = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
+					err = c.cli.Agent().UpdateTTLOpts("service:"+svc.ID, "pass", "pass", new(api.QueryOptions).WithContext(c.ctx))
+					if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+						_ = c.cli.Agent().ServiceDeregister(svc.ID)
+						return
+					}
 					if err != nil {
 						log.Errorf("[Consul] update ttl heartbeat to consul failed! err=%v", err)
 						// when the previous report fails, try to re register the service
diff --git a/transport/http/server.go b/transport/http/server.go
index c967e4a1..9c5550b2 100644
--- a/transport/http/server.go
+++ b/transport/http/server.go
@@ -141,6 +141,18 @@ func PathPrefix(prefix string) ServerOption {
 	}
 }
 
+func NotFoundHandler(handler http.Handler) ServerOption {
+	return func(s *Server) {
+		s.router.NotFoundHandler = handler
+	}
+}
+
+func MethodNotAllowedHandler(handler http.Handler) ServerOption {
+	return func(s *Server) {
+		s.router.MethodNotAllowedHandler = handler
+	}
+}
+
 // Server is an HTTP server wrapper.
 type Server struct {
 	*http.Server
@@ -177,12 +189,12 @@ func NewServer(opts ...ServerOption) *Server {
 		strictSlash: true,
 		router:      mux.NewRouter(),
 	}
+	srv.router.NotFoundHandler = http.DefaultServeMux
+	srv.router.MethodNotAllowedHandler = http.DefaultServeMux
 	for _, o := range opts {
 		o(srv)
 	}
 	srv.router.StrictSlash(srv.strictSlash)
-	srv.router.NotFoundHandler = http.DefaultServeMux
-	srv.router.MethodNotAllowedHandler = http.DefaultServeMux
 	srv.router.Use(srv.filter())
 	srv.Server = &http.Server{
 		Handler:   FilterChain(srv.filters...)(srv.router),


Test patch:
diff --git a/transport/http/server_test.go b/transport/http/server_test.go
index 781c06a7..44477504 100644
--- a/transport/http/server_test.go
+++ b/transport/http/server_test.go
@@ -371,3 +371,19 @@ func TestListener(t *testing.T) {
 		t.Errorf("expected not empty")
 	}
 }
+
+func TestNotFoundHandler(t *testing.T) {
+	mux := http.NewServeMux()
+	srv := NewServer(NotFoundHandler(mux))
+	if !reflect.DeepEqual(srv.router.NotFoundHandler, mux) {
+		t.Errorf("expected %v got %v", mux, srv.router.NotFoundHandler)
+	}
+}
+
+func TestMethodNotAllowedHandler(t *testing.T) {
+	mux := http.NewServeMux()
+	srv := NewServer(MethodNotAllowedHandler(mux))
+	if !reflect.DeepEqual(srv.router.MethodNotAllowedHandler, mux) {
+		t.Errorf("expected %v got %v", mux, srv.router.MethodNotAllowedHandler)
+	}
+}