seaweedfs/weed/operation/grpc_client.go

package operation

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"google.golang.org/grpc"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	connectionPool     = make(map[string]*util.ResourcePool)
	connectionPoolLock sync.Mutex
)

func WithVolumeServerClient(volumeServer string, grpcDialOption grpc.DialOption, fn func(context.Context, volume_server_pb.VolumeServerClient) error) error {

	ctx := context.Background()

	grpcAddress, err := toVolumeServerGrpcAddress(volumeServer)
	if err != nil {
		return err
	}

	return util.WithCachedGrpcClient(ctx, func(ctx2 context.Context, grpcConnection *grpc.ClientConn) error {
		client := volume_server_pb.NewVolumeServerClient(grpcConnection)
		return fn(ctx2, client)
	}, grpcAddress, grpcDialOption)

}

func toVolumeServerGrpcAddress(volumeServer string) (grpcAddress string, err error) {
	sepIndex := strings.LastIndex(volumeServer, ":")
	port, err := strconv.Atoi(volumeServer[sepIndex+1:])
	if err != nil {
		glog.Errorf("failed to parse volume server address: %v", volumeServer)
		return "", err
	}
	return fmt.Sprintf("%s:%d", volumeServer[0:sepIndex], port+10000), nil
}

func WithVolumeServerTcpConnection(volumeServer string, fn func(conn net.Conn) error) error {
	tcpAddress, err := toVolumeServerTcpAddress(volumeServer)
	if err != nil {
		return err
	}

	conn, err := getConnection(tcpAddress)
	if err != nil {
		return err
	}
	defer releaseConnection(conn, tcpAddress)

	err = fn(conn)
	return err
}

func getConnection(tcpAddress string) (net.Conn, error) {
	connectionPoolLock.Lock()
	defer connectionPoolLock.Unlock()

	pool, found := connectionPool[tcpAddress]
	if !found {
		println("creating pool for", tcpAddress)

		raddr, err := net.ResolveTCPAddr("tcp", tcpAddress)
		if err != nil {
			glog.Fatal(err)
		}

		pool = util.NewResourcePool(16, func() (interface{}, error) {

			conn, err := net.DialTCP("tcp", nil, raddr)
			if err != nil {
				return conn, err
			}
			conn.SetKeepAlive(true)
			conn.SetNoDelay(true)
			println("connected", tcpAddress, "=>", conn.LocalAddr().String())
			return conn, nil
		})
		connectionPool[tcpAddress] = pool
	}

	connObj, err := pool.Get(time.Minute)
	if err != nil {
		return nil, err
	}
	// println("get connection", tcpAddress, "=>", conn.LocalAddr().String())
	return connObj.(net.Conn), nil
}

func releaseConnection(conn net.Conn, tcpAddress string) {
	connectionPoolLock.Lock()
	defer connectionPoolLock.Unlock()

	pool, found := connectionPool[tcpAddress]
	if !found {
		println("can not return connection", tcpAddress, "=>", conn.LocalAddr().String())
		return
	}
	pool.Release(conn)
	// println("returned connection", tcpAddress, "=>", conn.LocalAddr().String())
}

func toVolumeServerTcpAddress(volumeServer string) (grpcAddress string, err error) {
	sepIndex := strings.LastIndex(volumeServer, ":")
	port, err := strconv.Atoi(volumeServer[sepIndex+1:])
	if err != nil {
		glog.Errorf("failed to parse volume server address: %v", volumeServer)
		return "", err
	}
	return fmt.Sprintf("%s:%d", volumeServer[0:sepIndex], port+20000), nil
}

func WithMasterServerClient(masterServer string, grpcDialOption grpc.DialOption, fn func(ctx2 context.Context, masterClient master_pb.SeaweedClient) error) error {

	ctx := context.Background()

	masterGrpcAddress, parseErr := util.ParseServerToGrpcAddress(masterServer)
	if parseErr != nil {
		return fmt.Errorf("failed to parse master grpc %v: %v", masterServer, parseErr)
	}

	return util.WithCachedGrpcClient(ctx, func(ctx2 context.Context, grpcConnection *grpc.ClientConn) error {
		client := master_pb.NewSeaweedClient(grpcConnection)
		return fn(ctx2, client)
	}, masterGrpcAddress, grpcDialOption)

}