operations/find-branches.go

package operations

import (
	"context"
	"io"

	"google.golang.org/grpc"

	pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
)

type fbn struct {
	repository *pb.Repository
	conn       *grpc.ClientConn
	client     pb.RefServiceClient
	f          func(fbn) error
}

// var GCFlagSet = flag.NewFlagSet("gc", flag.ExitOnError)
// var bitmapFlag = GCFlagSet.Bool("bitmap", false, "Create bitmap")

func (b fbn) Run() error {
	return b.f(b)
}

func findBranchNamesBenchmark(b fbn) error {
	request := &pb.FindAllBranchNamesRequest{
		Repository: b.repository,
	}

	stream, err := b.client.FindAllBranchNames(context.Background(), request)
	if err != nil {
		return err
	}

	for {
		_, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
}

func findBranchesBenchmark(b fbn) error {
	request := &pb.FindAllBranchesRequest{
		Repository: b.repository,
	}

	stream, err := b.client.FindAllBranches(context.Background(), request)
	if err != nil {
		return err
	}

	for {
		_, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
}

// NewFindAllBranchNames creates a new benchmark operation
func NewFindAllBranchNames(repository *pb.Repository, conn *grpc.ClientConn) BenchmarkOperation {
	client := pb.NewRefServiceClient(conn)

	return fbn{repository: repository, conn: conn, client: client, f: findBranchNamesBenchmark}
}

// NewFindAllBranches creates a new benchmark operation
func NewFindAllBranches(repository *pb.Repository, conn *grpc.ClientConn) BenchmarkOperation {
	client := pb.NewRefServiceClient(conn)

	return fbn{repository: repository, conn: conn, client: client, f: findBranchesBenchmark}
}