Refactor collector mains

cmd/all-in-one/main.go

@@ -19,44 +19,29 @@ import (
 	"fmt"
 	"io"
 	"log"
-	"net"
-	"net/http"
 	"os"
-	"strconv"
 
-	"github.com/gorilla/mux"
 	"github.com/opentracing/opentracing-go"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	jaegerClientConfig "github.com/uber/jaeger-client-go/config"
 	jaegerClientZapLog "github.com/uber/jaeger-client-go/log/zap"
 	"github.com/uber/jaeger-lib/metrics"
-	"github.com/uber/tchannel-go"
-	"github.com/uber/tchannel-go/thrift"
 	_ "go.uber.org/automaxprocs"
 	"go.uber.org/zap"
-	"google.golang.org/grpc"
 
 	agentApp "github.com/jaegertracing/jaeger/cmd/agent/app"
 	agentRep "github.com/jaegertracing/jaeger/cmd/agent/app/reporter"
 	agentGrpcRep "github.com/jaegertracing/jaeger/cmd/agent/app/reporter/grpc"
 	agentTchanRep "github.com/jaegertracing/jaeger/cmd/agent/app/reporter/tchannel"
 	"github.com/jaegertracing/jaeger/cmd/all-in-one/setupcontext"
 	collectorApp "github.com/jaegertracing/jaeger/cmd/collector/app"
-	collector "github.com/jaegertracing/jaeger/cmd/collector/app/builder"
-	"github.com/jaegertracing/jaeger/cmd/collector/app/grpcserver"
-	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling"
-	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
-	"github.com/jaegertracing/jaeger/cmd/collector/app/zipkin"
 	"github.com/jaegertracing/jaeger/cmd/docs"
 	"github.com/jaegertracing/jaeger/cmd/env"
 	"github.com/jaegertracing/jaeger/cmd/flags"
 	queryApp "github.com/jaegertracing/jaeger/cmd/query/app"
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
-	clientcfgHandler "github.com/jaegertracing/jaeger/pkg/clientcfg/clientcfghttp"
 	"github.com/jaegertracing/jaeger/pkg/config"
-	"github.com/jaegertracing/jaeger/pkg/healthcheck"
-	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
 	"github.com/jaegertracing/jaeger/pkg/version"
 	ss "github.com/jaegertracing/jaeger/plugin/sampling/strategystore"
 	"github.com/jaegertracing/jaeger/plugin/storage"
@@ -65,9 +50,6 @@ import (
 	"github.com/jaegertracing/jaeger/storage/dependencystore"
 	"github.com/jaegertracing/jaeger/storage/spanstore"
 	storageMetrics "github.com/jaegertracing/jaeger/storage/spanstore/metrics"
-	jc "github.com/jaegertracing/jaeger/thrift-gen/jaeger"
-	sc "github.com/jaegertracing/jaeger/thrift-gen/sampling"
-	zc "github.com/jaegertracing/jaeger/thrift-gen/zipkincore"
 )
 
 // all-in-one/main is a standalone full-stack jaeger backend, backed by a memory store
@@ -108,6 +90,7 @@ by default uses only in-memory database.`,
 			if err := storageFactory.Initialize(metricsFactory, logger); err != nil {
 				logger.Fatal("Failed to init storage factory", zap.Error(err))
 			}
+
 			spanReader, err := storageFactory.CreateSpanReader()
 			if err != nil {
 				logger.Fatal("Failed to create span reader", zap.Error(err))
@@ -122,25 +105,40 @@ by default uses only in-memory database.`,
 			}
 
 			strategyStoreFactory.InitFromViper(v)
-			strategyStore := initSamplingStrategyStore(strategyStoreFactory, metricsFactory, logger)
+			if err := strategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
+				logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
+			}
+			strategyStore, err := strategyStoreFactory.CreateStrategyStore()
+			if err != nil {
+				logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
+			}
 
 			aOpts := new(agentApp.Builder).InitFromViper(v)
 			repOpts := new(agentRep.Options).InitFromViper(v, logger)
 			tchanBuilder := agentTchanRep.NewBuilder().InitFromViper(v, logger)
 			grpcBuilder := agentGrpcRep.NewConnBuilder().InitFromViper(v)
-			cOpts := new(collector.CollectorOptions).InitFromViper(v)
+			cOpts := new(collectorApp.CollectorOptions).InitFromViper(v)
 			qOpts := new(queryApp.QueryOptions).InitFromViper(v, logger)
 
-			collectorSrv := startCollector(cOpts, spanWriter, logger, metricsFactory, strategyStore, svc.HC())
-			startAgent(aOpts, repOpts, tchanBuilder, grpcBuilder, cOpts, logger, metricsFactory)
+			c := collectorApp.New(&collectorApp.CollectorParams{
+				ServiceName:    "jaeger-collector",
+				Logger:         logger,
+				MetricsFactory: metricsFactory,
+				SpanWriter:     spanWriter,
+				StrategyStore:  strategyStore,
+				HealthCheck:    svc.HC(),
+			})
+			c.Start(cOpts)
+
+			startAgent(aOpts, repOpts, tchanBuilder, grpcBuilder, cOpts.CollectorGRPCPort, logger, metricsFactory)
 			querySrv := startQuery(
 				svc, qOpts, archiveOptions(storageFactory, logger),
 				spanReader, dependencyReader,
 				rootMetricsFactory, metricsFactory,
 			)
 
 			svc.RunAndThen(func() {
-				collectorSrv.GracefulStop()
+				c.Close()
 				querySrv.Close()
 				if closer, ok := spanWriter.(io.Closer); ok {
 					err := closer.Close()
@@ -167,7 +165,7 @@ by default uses only in-memory database.`,
 		agentRep.AddFlags,
 		agentTchanRep.AddFlags,
 		agentGrpcRep.AddFlags,
-		collector.AddFlags,
+		collectorApp.AddFlags,
 		queryApp.AddFlags,
 		strategyStoreFactory.AddFlags,
 	)
@@ -183,13 +181,13 @@ func startAgent(
 	repOpts *agentRep.Options,
 	tchanBuilder *agentTchanRep.Builder,
 	grpcBuilder *agentGrpcRep.ConnBuilder,
-	cOpts *collector.CollectorOptions,
+	collectorGRPCPort int,
 	logger *zap.Logger,
 	baseFactory metrics.Factory,
 ) {
 	metricsFactory := baseFactory.Namespace(metrics.NSOptions{Name: "agent", Tags: nil})
 
-	grpcBuilder.CollectorHostPorts = append(grpcBuilder.CollectorHostPorts, fmt.Sprintf("127.0.0.1:%d", cOpts.CollectorGRPCPort))
+	grpcBuilder.CollectorHostPorts = append(grpcBuilder.CollectorHostPorts, fmt.Sprintf("127.0.0.1:%d", collectorGRPCPort))
 	cp, err := agentApp.CreateCollectorProxy(repOpts, tchanBuilder, grpcBuilder, logger, metricsFactory)
 	if err != nil {
 		logger.Fatal("Could not create collector proxy", zap.Error(err))
@@ -206,116 +204,6 @@ func startAgent(
 	}
 }
 
-func startCollector(
-	cOpts *collector.CollectorOptions,
-	spanWriter spanstore.Writer,
-	logger *zap.Logger,
-	baseFactory metrics.Factory,
-	strategyStore strategystore.StrategyStore,
-	hc *healthcheck.HealthCheck,
-) *grpc.Server {
-	metricsFactory := baseFactory.Namespace(metrics.NSOptions{Name: "collector", Tags: nil})
-
-	spanHandlerBuilder := &collector.SpanHandlerBuilder{
-		SpanWriter:     spanWriter,
-		CollectorOpts:  *cOpts,
-		Logger:         logger,
-		MetricsFactory: metricsFactory,
-	}
-
-	zipkinSpansHandler, jaegerBatchesHandler, grpcHandler := spanHandlerBuilder.BuildHandlers()
-
-	{
-		ch, err := tchannel.NewChannel("jaeger-collector", &tchannel.ChannelOptions{})
-		if err != nil {
-			logger.Fatal("Unable to create new TChannel", zap.Error(err))
-		}
-		server := thrift.NewServer(ch)
-		batchHandler := collectorApp.NewTChannelHandler(jaegerBatchesHandler, zipkinSpansHandler)
-		server.Register(jc.NewTChanCollectorServer(batchHandler))
-		server.Register(zc.NewTChanZipkinCollectorServer(batchHandler))
-		server.Register(sc.NewTChanSamplingManagerServer(sampling.NewHandler(strategyStore)))
-		portStr := ":" + strconv.Itoa(cOpts.CollectorPort)
-		listener, err := net.Listen("tcp", portStr)
-		if err != nil {
-			logger.Fatal("Unable to start listening on channel", zap.Error(err))
-		}
-		logger.Info("Starting jaeger-collector TChannel server", zap.Int("port", cOpts.CollectorPort))
-		logger.Warn("TChannel has been deprecated and will be removed in a future release")
-		ch.Serve(listener)
-	}
-
-	server, err := startGRPCServer(cOpts.CollectorGRPCPort, grpcHandler, strategyStore, logger)
-	if err != nil {
-		logger.Fatal("Could not start gRPC collector", zap.Error(err))
-	}
-
-	{
-		r := mux.NewRouter()
-		apiHandler := collectorApp.NewAPIHandler(jaegerBatchesHandler)
-		apiHandler.RegisterRoutes(r)
-
-		cfgHandler := clientcfgHandler.NewHTTPHandler(clientcfgHandler.HTTPHandlerParams{
-			ConfigManager: &clientcfgHandler.ConfigManager{
-				SamplingStrategyStore: strategyStore,
-				// TODO provide baggage manager
-			},
-			MetricsFactory:         metricsFactory,
-			BasePath:               "/api",
-			LegacySamplingEndpoint: false,
-		})
-		cfgHandler.RegisterRoutes(r)
-
-		recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)
-		go startZipkinHTTPAPI(logger, cOpts.CollectorZipkinHTTPPort, zipkinSpansHandler, recoveryHandler)
-
-		httpPortStr := ":" + strconv.Itoa(cOpts.CollectorHTTPPort)
-		logger.Info("Starting jaeger-collector HTTP server", zap.String("http-host-port", httpPortStr))
-		go func() {
-			if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
-				logger.Fatal("Could not launch jaeger-collector HTTP server", zap.Error(err))
-			}
-			hc.Set(healthcheck.Unavailable)
-		}()
-	}
-	return server
-}
-
-func startGRPCServer(
-	port int,
-	handler *collectorApp.GRPCHandler,
-	samplingStore strategystore.StrategyStore,
-	logger *zap.Logger,
-) (*grpc.Server, error) {
-	server := grpc.NewServer()
-	_, err := grpcserver.StartGRPCCollector(port, server, handler, samplingStore, logger, func(err error) {
-		logger.Fatal("gRPC collector failed", zap.Error(err))
-	})
-	if err != nil {
-		return nil, err
-	}
-	return server, err
-}
-
-func startZipkinHTTPAPI(
-	logger *zap.Logger,
-	zipkinPort int,
-	zipkinSpansHandler collectorApp.ZipkinSpansHandler,
-	recoveryHandler func(http.Handler) http.Handler,
-) {
-	if zipkinPort != 0 {
-		r := mux.NewRouter()
-		zHandler := zipkin.NewAPIHandler(zipkinSpansHandler)
-		zHandler.RegisterRoutes(r)
-		httpPortStr := ":" + strconv.Itoa(zipkinPort)
-		logger.Info("Listening for Zipkin HTTP traffic", zap.Int("zipkin.http-port", zipkinPort))
-
-		if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
-			logger.Fatal("Could not launch service", zap.Error(err))
-		}
-	}
-}
-
 func startQuery(
 	svc *flags.Service,
 	qOpts *queryApp.QueryOptions,
@@ -334,21 +222,6 @@ func startQuery(
 	return server
 }
 
-func initSamplingStrategyStore(
-	samplingStrategyStoreFactory *ss.Factory,
-	metricsFactory metrics.Factory,
-	logger *zap.Logger,
-) strategystore.StrategyStore {
-	if err := samplingStrategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
-		logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
-	}
-	strategyStore, err := samplingStrategyStoreFactory.CreateStrategyStore()
-	if err != nil {
-		logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
-	}
-	return strategyStore
-}
-
 func archiveOptions(storageFactory istorage.Factory, logger *zap.Logger) *querysvc.QueryServiceOptions {
 	opts := &querysvc.QueryServiceOptions{}
 	if !opts.InitArchiveStorage(storageFactory, logger) {

cmd/collector/app/builder/builder_flags.go → cmd/collector/app/builder_flags.go

@@ -13,14 +13,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package builder
+package app
 
 import (
 	"flag"
 
 	"github.com/spf13/viper"
 
-	"github.com/jaegertracing/jaeger/cmd/collector/app"
 	"github.com/jaegertracing/jaeger/cmd/flags"
 	"github.com/jaegertracing/jaeger/pkg/config/tlscfg"
 	"github.com/jaegertracing/jaeger/ports"
@@ -74,8 +73,8 @@ type CollectorOptions struct {
 // AddFlags adds flags for CollectorOptions
 func AddFlags(flags *flag.FlagSet) {
 	flags.Uint(collectorDynQueueSizeMemory, 0, "(experimental) The max memory size in MiB to use for the dynamic queue.")
-	flags.Int(collectorQueueSize, app.DefaultQueueSize, "The queue size of the collector")
-	flags.Int(collectorNumWorkers, app.DefaultNumWorkers, "The number of workers pulling items from the queue")
+	flags.Int(collectorQueueSize, DefaultQueueSize, "The queue size of the collector")
+	flags.Int(collectorNumWorkers, DefaultNumWorkers, "The number of workers pulling items from the queue")
 	flags.Int(collectorPort, ports.CollectorTChannel, "The TChannel port for the collector service")
 	flags.Int(collectorHTTPPort, ports.CollectorHTTP, "The HTTP port for the collector service")
 	flags.Int(collectorGRPCPort, ports.CollectorGRPC, "The gRPC port for the collector service")