Full MVC of gitlab-zoekt-indexer

cmd/gitlab-zoekt-indexer/main.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"io"
 	"log"
 	"net/http"
 	"os"
@@ -14,20 +13,22 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/go-chi/chi/v5"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"gitlab.com/gitlab-org/search-team/gitlab-zoekt-indexer/correlation"
 	"gitlab.com/gitlab-org/search-team/gitlab-zoekt-indexer/indexer"
 )
 
-type Options struct {
+type options struct {
 	indexDir string
 	listen   string
 }
 
 type indexServer struct {
-	opts                 Options
+	indexDir             string
 	promRegistry         *prometheus.Registry
 	metricsRequestsTotal *prometheus.CounterVec
 }
@@ -42,145 +43,217 @@ type gitalyConnectionInfo struct {
 type indexRequest struct {
 	Timeout              string
 	RepoID               uint32
-	CorrelationID        string
 	GitalyConnectionInfo *gitalyConnectionInfo
 	FileSizeLimit        int
 }
 
-type statusRequest struct {
-	RepoID uint32
-}
+func (s *indexServer) indexRepository(req indexRequest, ctx context.Context) error {
+	timeout, err := time.ParseDuration(req.Timeout)
+	if err != nil {
+		return fmt.Errorf("failed to parse Timeout: %v with error %v", req.Timeout, err)
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
 
-func (o *Options) createIndexDir() {
-	if err := os.MkdirAll(o.indexDir, 0o755); err != nil {
-		log.Fatalf("createIndexDir %s: %v", o.indexDir, err)
+	idx := &indexer.Indexer{
+		IndexDir:           s.indexDir,
+		ProjectID:          req.RepoID,
+		CorrelationID:      correlation.GetCorrelationID(ctx),
+		GitalyAddress:      req.GitalyConnectionInfo.Address,
+		GitalyStorageName:  req.GitalyConnectionInfo.Storage,
+		GitalyRelativePath: req.GitalyConnectionInfo.Path,
+		LimitFileSize:      req.FileSizeLimit,
 	}
+
+	if err := idx.IndexRepository(ctx); err != nil {
+		return err
+	}
+
+	return nil
 }
 
-func (s *indexServer) serveHealthCheck(w http.ResponseWriter, r *http.Request) {
-	// Nothing to do. Just return 200
+func (s *indexServer) createIndexDir() {
+	if err := os.MkdirAll(s.indexDir, 0o755); err != nil {
+		log.Fatalf("createIndexDir %s: %v", s.indexDir, err)
+	}
 }
 
-func (s *indexServer) serveStatus(w http.ResponseWriter, r *http.Request) {
+func (s *indexServer) handleStatus() http.HandlerFunc {
 	route := "status"
 
-	dec := json.NewDecoder(r.Body)
-	dec.DisallowUnknownFields()
-	var req statusRequest
-	err := dec.Decode(&req)
-
-	if err != nil {
-		http.Error(w, "JSON parser error", http.StatusBadRequest)
-		return
+	type response struct {
+		Success bool
+		SHA     string
 	}
 
-	idx := &indexer.Indexer{
-		IndexDir:  s.opts.indexDir,
-		ProjectID: req.RepoID,
-	}
+	return func(w http.ResponseWriter, r *http.Request) {
+		param := chi.URLParam(r, "id")
+		repoID, err := strconv.ParseUint(param, 10, 32)
 
-	currentSHA, err := idx.CurrentSHA()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
 
-	if err != nil {
-		s.respondWithError(w, r.Method, route, err)
-		return
-	}
+		idx := &indexer.Indexer{
+			IndexDir:  s.indexDir,
+			ProjectID: uint32(repoID),
+		}
 
-	response := map[string]any{
-		"Success": true,
-		"SHA":     currentSHA,
-	}
+		currentSHA, err := idx.CurrentSHA()
 
-	if err != nil {
-		s.respondWithError(w, r.Method, route, err)
-		return
-	}
+		if err != nil {
+			s.respondWithError(w, r, route, err)
+			return
+		}
 
-	w.Header().Set("Content-Type", "application/json")
-	_ = json.NewEncoder(w).Encode(response)
+		if err != nil {
+			s.respondWithError(w, r, route, err)
+			return
+		}
 
-	s.incrementRequestsTotal(r.Method, route, http.StatusOK)
+		resp := response{
+			Success: true,
+			SHA:     currentSHA,
+		}
+
+		s.respondWith(w, r, route, resp)
+	}
 }
 
-func (s *indexServer) serveMetrics(w http.ResponseWriter, r *http.Request) {
-	promhttp.HandlerFor(s.promRegistry, promhttp.HandlerOpts{Registry: s.promRegistry}).ServeHTTP(w, r)
+func (s *indexServer) handleMetrics() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		promhttp.HandlerFor(s.promRegistry, promhttp.HandlerOpts{Registry: s.promRegistry}).ServeHTTP(w, r)
+	}
 }
 
-func parseRequest(r io.Reader) (indexRequest, error) {
-	dec := json.NewDecoder(r)
+func (s *indexServer) decode(r *http.Request, v interface{}) error {
+	dec := json.NewDecoder(r.Body)
 	dec.DisallowUnknownFields()
-	var req indexRequest
-	err := dec.Decode(&req)
+	return dec.Decode(v)
+}
 
-	if err != nil {
-		return req, errors.New("json parser error")
-	}
+func (s *indexServer) handleIndex() http.HandlerFunc {
+	route := "index"
 
-	if req.GitalyConnectionInfo == nil {
-		return req, errors.New("gitalyConnectionInfo Error")
+	type response struct {
+		Success bool
 	}
 
-	if req.FileSizeLimit == 0 {
-		return req, errors.New("fileSizeLimit is not set")
-	}
+	parseRequest := func(r *http.Request) (indexRequest, error) {
+		var req indexRequest
+		err := s.decode(r, &req)
 
-	return req, nil
-}
+		if err != nil {
+			return req, errors.New("json parser error")
+		}
 
-func (s *indexServer) serveIndex(w http.ResponseWriter, r *http.Request) {
-	route := "index"
+		if req.GitalyConnectionInfo == nil {
+			return req, errors.New("gitalyConnectionInfo Error")
+		}
 
-	req, err := parseRequest(r.Body)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
+		if req.FileSizeLimit == 0 {
+			return req, errors.New("fileSizeLimit is not set")
+		}
 
-	response, err := s.indexRepository(req)
-	if err != nil {
-		s.respondWithError(w, r.Method, route, err)
-		return
+		return req, nil
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-	_ = json.NewEncoder(w).Encode(response)
+	return func(w http.ResponseWriter, r *http.Request) {
+		req, err := parseRequest(r)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
 
-	s.incrementRequestsTotal(r.Method, route, http.StatusOK)
+		err = s.indexRepository(req, r.Context())
+		if err != nil {
+			s.respondWithError(w, r, route, err)
+			return
+		}
+
+		resp := response{
+			Success: true,
+		}
+
+		s.respondWith(w, r, route, resp)
+	}
 }
 
-func (s *indexServer) serveTruncate(w http.ResponseWriter, r *http.Request) {
+func (s *indexServer) handleTruncate() http.HandlerFunc {
 	route := "truncate"
-	err := emptyDirectory(s.opts.indexDir)
 
-	if err != nil {
-		err = fmt.Errorf("failed to empty indexDir: %v with error: %v", s.opts.indexDir, err)
+	type response struct {
+		Success bool
+	}
 
-		s.respondWithError(w, r.Method, route, err)
-		return
+	emptyDirectory := func(dir string) error {
+		files, err := os.ReadDir(dir)
+
+		if err != nil {
+			return err
+		}
+
+		for _, file := range files {
+			filePath := filepath.Join(dir, file.Name())
+			err := os.RemoveAll(filePath)
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
 	}
 
-	response := map[string]any{
-		"Success": true,
+	return func(w http.ResponseWriter, r *http.Request) {
+		err := emptyDirectory(s.indexDir)
+
+		if err != nil {
+			err = fmt.Errorf("failed to empty indexDir: %v with error: %v", s.indexDir, err)
+
+			s.respondWithError(w, r, route, err)
+			return
+		}
+
+		resp := response{
+			Success: true,
+		}
+
+		s.respondWith(w, r, route, resp)
 	}
+}
+
+func (s *indexServer) respondWith(w http.ResponseWriter, r *http.Request, route string, data interface{}) {
 	w.Header().Set("Content-Type", "application/json")
-	_ = json.NewEncoder(w).Encode(response)
+
+	if err := json.NewEncoder(w).Encode(data); err != nil {
+		s.respondWithError(w, r, route, err)
+		return
+	}
 
 	s.incrementRequestsTotal(r.Method, route, http.StatusOK)
 }
 
-func (s *indexServer) respondWithError(w http.ResponseWriter, method, route string, err error) {
-	responseCode := http.StatusInternalServerError
+func (s *indexServer) respondWithError(w http.ResponseWriter, r *http.Request, route string, err error) {
+	type response struct {
+		Success bool
+		Error   string
+	}
 
-	s.incrementRequestsTotal(method, route, responseCode)
+	responseCode := http.StatusInternalServerError
 
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(responseCode)
-	response := map[string]any{
-		"Success": false,
-		"Error":   err.Error(),
+
+	resp := response{
+		Success: false,
+		Error:   err.Error(),
 	}
 
-	_ = json.NewEncoder(w).Encode(response)
+	_ = json.NewEncoder(w).Encode(resp)
+
+	s.incrementRequestsTotal(r.Method, route, responseCode)
 }
 
 func (s *indexServer) incrementRequestsTotal(method, route string, responseCode int) {
@@ -205,70 +278,7 @@ func (s *indexServer) initMetrics() {
 	)
 }
 
-func (s *indexServer) startIndexingApi() {
-	s.initMetrics()
-	s.opts.createIndexDir()
-
-	http.HandleFunc("/", s.serveHealthCheck)
-	http.HandleFunc("/status", s.serveStatus)
-	http.HandleFunc("/index", s.serveIndex)
-	http.HandleFunc("/truncate", s.serveTruncate)
-	http.HandleFunc("/metrics", s.serveMetrics)
-
-	if err := http.ListenAndServe(s.opts.listen, nil); err != nil {
-		log.Fatal(err)
-	}
-}
-
-func emptyDirectory(dir string) error {
-	files, err := os.ReadDir(dir)
-
-	if err != nil {
-		return err
-	}
-
-	for _, file := range files {
-		filePath := filepath.Join(dir, file.Name())
-		err := os.RemoveAll(filePath)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (s *indexServer) indexRepository(req indexRequest) (map[string]any, error) {
-	timeout, err := time.ParseDuration(req.Timeout)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse Timeout: %v with error %v", req.Timeout, err)
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	defer cancel()
-
-	idx := &indexer.Indexer{
-		IndexDir:           s.opts.indexDir,
-		ProjectID:          req.RepoID,
-		CorrelationID:      req.CorrelationID,
-		GitalyAddress:      req.GitalyConnectionInfo.Address,
-		GitalyStorageName:  req.GitalyConnectionInfo.Storage,
-		GitalyRelativePath: req.GitalyConnectionInfo.Path,
-		LimitFileSize:      req.FileSizeLimit,
-	}
-
-	if err := idx.IndexRepository(ctx); err != nil {
-		return nil, err
-	}
-
-	response := map[string]any{
-		"Success": true,
-	}
-
-	return response, nil
-}
-
-func parseOptions() Options {
+func parseOptions() options {
 	indexDir := flag.String("index_dir", "", "directory holding index shards.")
 	listen := flag.String("listen", ":6060", "listen on this address.")
 	flag.Parse()
@@ -277,18 +287,34 @@ func parseOptions() Options {
 		log.Fatal("must set -index_dir")
 	}
 
-	return Options{
+	return options{
 		indexDir: *indexDir,
 		listen:   *listen,
 	}
 }
 
 func main() {
+	opts := parseOptions()
+
 	server := indexServer{
-		opts: parseOptions(),
+		indexDir: opts.indexDir,
 	}
 
-	log.Printf("Starting server on %s", server.opts.listen)
+	server.startIndexingApi(opts.listen)
+}
 
-	server.startIndexingApi()
+func (s *indexServer) startIndexingApi(listen string) {
+	s.initMetrics()
+	s.createIndexDir()
+
+	httpServer := http.Server{
+		Addr:    listen,
+		Handler: s.router(),
+	}
+
+	log.Printf("Starting server on %s", listen)
+
+	if err := httpServer.ListenAndServe(); err != nil {
+		log.Fatal(err)
+	}
 }