Repository: go-kratos/kratos
Pull request #3146 (2024-01-04 08:55:04)

fix: clear err after app stop err handling
#### Description (what this PR does / why we need it):
clear the err after the app wait err handing, to avoid the err return when afterStop is not set.


#### Which issue(s) this PR fixes (resolves / be part of):
<!--
* Automatically closes linked issue when PR is merged.
* If your PR is not fully resolved the issue, please use `part of #<issue number>` instead.

Usage: `fixes/resolves #<issue number>`, or `fixes/resolves (paste link of issue)`.
-->


#### Other special notes for the reviewers:
<!--
* Somethings that need extra attention for the reviewers
* Some additional notes, TODO list, etc.
-->



Patch:
diff --git a/app.go b/app.go
index 2fa52ec9..1fa23132 100644
--- a/app.go
+++ b/app.go
@@ -139,6 +139,7 @@ func (a *App) Run() error {
 	if err = eg.Wait(); err != nil && !errors.Is(err, context.Canceled) {
 		return err
 	}
+	err = nil
 	for _, fn := range a.opts.afterStop {
 		err = fn(sctx)
 	}
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)
+	}
+}