pkg/cmdsetup/data.go

// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package cmdsetup

import (
	"context"

	"github.com/spf13/cobra"

	"github.com/apache/skywalking-banyandb/api/common"
	"github.com/apache/skywalking-banyandb/banyand/liaison/grpc/route"
	"github.com/apache/skywalking-banyandb/banyand/measure"
	"github.com/apache/skywalking-banyandb/banyand/metadata/service"
	"github.com/apache/skywalking-banyandb/banyand/observability"
	"github.com/apache/skywalking-banyandb/banyand/observability/services"
	"github.com/apache/skywalking-banyandb/banyand/property"
	"github.com/apache/skywalking-banyandb/banyand/protector"
	"github.com/apache/skywalking-banyandb/banyand/query"
	"github.com/apache/skywalking-banyandb/banyand/queue"
	"github.com/apache/skywalking-banyandb/banyand/queue/sub"
	"github.com/apache/skywalking-banyandb/banyand/stream"
	"github.com/apache/skywalking-banyandb/banyand/trace"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/run"
	"github.com/apache/skywalking-banyandb/pkg/version"
)

func newDataCmd(runners ...run.Unit) *cobra.Command {
	l := logger.GetLogger("bootstrap")
	ctx := context.Background()
	metaSvc, err := service.NewService()
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate metadata service")
	}
	metricsPipeline := queue.Local()
	metricSvc := services.NewMetricService(metaSvc, metricsPipeline, "data", nil)
	metaSvc.SetMetricsRegistry(metricSvc)
	pm := protector.NewMemory(metricSvc)
	pipeline := sub.NewServer(metricSvc)
	pipeline.SetNodeSchemaStatusRepo(metaSvc)
	metaSvc.SetSnapshotPipeline(pipeline)
	propertyStreamPipeline := queue.Local()
	metaSvc.SetPropertyPipelineClient(propertyStreamPipeline)
	propertySvc, err := property.NewService(metaSvc, pipeline, propertyStreamPipeline, metricSvc, pm)
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate property service")
	}
	pipeline.SetRouteProviders(map[string]route.TableProvider{
		"property": propertySvc,
	})

	streamSvc, err := stream.NewService(metaSvc, pipeline, metricSvc, pm, propertyStreamPipeline)
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate stream service")
	}
	measureSvc, err := measure.NewDataSVC(metaSvc, pipeline, metricsPipeline, metricSvc, pm)
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate measure service")
	}
	traceSvc, err := trace.NewService(metaSvc, pipeline, metricSvc, pm)
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate trace service")
	}
	q, err := query.NewService(ctx, streamSvc, measureSvc, traceSvc, metaSvc, pipeline, metricSvc, true)
	if err != nil {
		l.Fatal().Err(err).Msg("failed to initiate query processor")
	}
	profSvc := observability.NewProfService()

	var units []run.Unit
	units = append(units, runners...)
	units = append(units,
		metricsPipeline,
		metricSvc,
		metaSvc,
		pm,
		pipeline,
		propertyStreamPipeline,
		propertySvc,
		measureSvc,
		streamSvc,
		traceSvc,
		q,
		profSvc,
	)
	dataGroup := run.NewGroup("data")
	dataGroup.Register(units...)
	dataCmd := &cobra.Command{
		Use:     "data",
		Version: version.Build(),
		Short:   "Run as the data server",
		RunE: func(_ *cobra.Command, _ []string) (err error) {
			// See standalone.go for the rationale: Cobra short-circuits
			// PersistentPostRunE on a non-nil RunE error, so we must
			// drain the artifact sink here.
			defer ShutdownSupervisor()
			node, err := common.GenerateNode(pipeline.GetPort(), nil, propertySvc.GetGossIPGrpcPort(), metaSvc.GetSchemaServerPort(), metaSvc.GetSchemaGossipPort())
			if err != nil {
				return err
			}
			logger.GetLogger().Info().Msg("starting as a data server")
			// Spawn our go routines and wait for shutdown. The Run context is
			// derived from SupervisorContext so that a panic recovered anywhere
			// in the process, including goroutines spawned via run.Go, fires
			// cancellation here.
			if err := dataGroup.Run(context.WithValue(SupervisorContext(), common.ContextNodeKey, node)); err != nil {
				logger.GetLogger().Error().Err(err).Stack().Str("name", dataGroup.Name()).Msg("Exit")
				return err
			}
			return nil
		},
	}
	dataCmd.Flags().AddFlagSet(dataGroup.RegisterFlags().FlagSet)
	return dataCmd
}