Return errors when http/grpc servers exit.

There was a case where the HTTP server died but the process exited. We
need to exit when the servers exit.

Signed-off-by: Goutham Veeramachaneni <cs14btech11014@iith.ac.in>
Signed-off-by: Goutham Veeramachaneni <gouthamve@gmail.com>
Signed-off-by: Goutham Veeramachaneni <cs14btech11014@iith.ac.in>

server/server.go

@@ -160,18 +160,48 @@ func RegisterInstrumentation(router *mux.Router) {
 	router.PathPrefix("/debug/pprof").Handler(http.DefaultServeMux)
 }
 
-// Run the server; blocks until SIGTERM is received.
-func (s *Server) Run() {
-	go s.httpServer.Serve(s.httpListener)
+// Run the server; blocks until SIGTERM or an error is received.
+func (s *Server) Run() error {
+	errChan := make(chan error)
+
+	go func() {
+		err := s.httpServer.Serve(s.httpListener)
+		if err == http.ErrServerClosed {
+			err = nil
+		}
+
+		select {
+		case errChan <- err:
+		default:
+		}
+	}()
 
 	// Setup gRPC server
 	// for HTTP over gRPC, ensure we don't double-count the middleware
 	httpgrpc.RegisterHTTPServer(s.GRPC, httpgrpc_server.NewServer(s.HTTP))
-	go s.GRPC.Serve(s.grpcListener)
+	go func() {
+		err := s.GRPC.Serve(s.grpcListener)
+		if err == grpc.ErrServerStopped {
+			err = nil
+		}
+
+		select {
+		case errChan <- err:
+		default:
+		}
+	}()
 	defer s.GRPC.GracefulStop()
 
 	// Wait for a signal
-	s.handler.Loop()
+	go func() {
+		s.handler.Loop()
+		select {
+		case errChan <- nil:
+		default:
+		}
+	}()
+
+	return <-errChan
 }
 
 // Stop unblocks Run().

server/server_test.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"strconv"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/status"
@@ -94,3 +95,41 @@ func TestErrorInstrumentationMiddleware(t *testing.T) {
 		"/server.FakeServer/Succeed":           "success",
 	}, statuses)
 }
+
+func TestRunReturnsError(t *testing.T) {
+	cfg := Config{
+		HTTPListenPort: 9190,
+		GRPCListenPort: 9191,
+	}
+	t.Run("http", func(t *testing.T) {
+		cfg.MetricsNamespace = "testing_http"
+		srv, err := New(cfg)
+		require.NoError(t, err)
+
+		errChan := make(chan error, 1)
+		go func() {
+			errChan <- srv.Run()
+		}()
+
+		time.Sleep(1 * time.Second) // For all go-routines to initialise.
+
+		require.NoError(t, srv.httpListener.Close())
+		require.NotNil(t, <-errChan)
+	})
+
+	t.Run("grpc", func(t *testing.T) {
+		cfg.MetricsNamespace = "testing_grpc"
+		srv, err := New(cfg)
+		require.NoError(t, err)
+
+		errChan := make(chan error, 1)
+		go func() {
+			errChan <- srv.Run()
+		}()
+
+		time.Sleep(1 * time.Second) // For all go-routines to initialise.
+
+		require.NoError(t, srv.grpcListener.Close())
+		require.NotNil(t, <-errChan)
+	})
+}