client.go

// Copyright 2019 The Cockroach Authors.
//
// Use of this software is governed by the CockroachDB Software License
// included in the /LICENSE file.

package blobs

import (
"context"
"io"

"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/blobs/blobspb"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/rpc"
"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
"github.com/cockroachdb/cockroach/pkg/util/ioctx"
"github.com/cockroachdb/errors"
"google.golang.org/grpc/metadata"
)

// BlobClient provides an interface for file access on all nodes' local storage.
// Given the nodeID of the node on which the operation should occur, the a blob
// client should be able to find the correct node and call its blob service API.
typeBlobClientinterface {
// ReadFile fetches the named payload from the requested node,
// and stores it in memory. It then returns an io.ReadCloser to
// read the contents.
ReadFile(ctx context.Context, filestring, offsetint64) (ioctx.ReadCloserCtx, int64, error)

// Writer opens the named payload on the requested node for writing.
Writer(ctx context.Context, filestring) (io.WriteCloser, error)

// List lists the corresponding filenames from the requested node.
// The requested node can be the current node.
List(ctx context.Context, patternstring) ([]string, error)

// Delete deletes the specified file or empty directory from a remote node.
Delete(ctx context.Context, filestring) error

// Stat gets the size (in bytes) of a specified file from a remote node.
Stat(ctx context.Context, filestring) (*blobspb.BlobStat, error)
}

var_BlobClient=&remoteClient{}

// remoteClient uses the node dialer and blob service clients
// to Read or Write bulk files from/to other nodes.
typeremoteClientstruct {
blobClient blobspb.BlobClient
}

// newRemoteClient instantiates a remote blob service client.
funcnewRemoteClient(blobClient blobspb.BlobClient) BlobClient {
return&remoteClient{blobClient: blobClient}
}

func (c*remoteClient) ReadFile(
ctx context.Context, filestring, offsetint64,
) (ioctx.ReadCloserCtx, int64, error) {
// Check that file exists before reading from it and get size to return.
st, err:=c.Stat(ctx, file)
iferr!=nil {
returnnil, 0, err
 }
stream, err:=c.blobClient.GetStream(ctx, &blobspb.GetRequest{
Filename: file,
Offset: offset,
 })
returnnewGetStreamReader(stream), st.Filesize, errors.Wrap(err, "fetching file")
}

typestreamWriterstruct {
s blobspb.Blob_PutStreamClient
buf blobspb.StreamChunk
}

func (w*streamWriter) Write(p []byte) (int, error) {
n:=0
forlen(p) >0 {
l:=copy(w.buf.Payload[:cap(w.buf.Payload)], p)
w.buf.Payload=w.buf.Payload[:l]
p=p[l:]
ifl>0 {
iferr:=w.s.Send(&w.buf); err!=nil {
returnn, err
 }
 }
n+=l
 }
returnn, nil
}

func (w*streamWriter) Close() error {
_, err:=w.s.CloseAndRecv()
returnerr
}

func (c*remoteClient) Writer(ctx context.Context, filestring) (io.WriteCloser, error) {
ctx=metadata.AppendToOutgoingContext(ctx, "filename", file)
stream, err:=c.blobClient.PutStream(ctx)
iferr!=nil {
returnnil, err
 }
buf:=make([]byte, 0, chunkSize)
return&streamWriter{s: stream, buf: blobspb.StreamChunk{Payload: buf}}, nil
}

func (c*remoteClient) List(ctx context.Context, patternstring) ([]string, error) {
resp, err:=c.blobClient.List(ctx, &blobspb.GlobRequest{
Pattern: pattern,
 })
iferr!=nil {
returnnil, errors.Wrap(err, "fetching list")
 }
returnresp.Files, nil
}

func (c*remoteClient) Delete(ctx context.Context, filestring) error {
_, err:=c.blobClient.Delete(ctx, &blobspb.DeleteRequest{
Filename: file,
 })
returnerr
}

func (c*remoteClient) Stat(ctx context.Context, filestring) (*blobspb.BlobStat, error) {
resp, err:=c.blobClient.Stat(ctx, &blobspb.StatRequest{
Filename: file,
 })
iferr!=nil {
returnnil, err
 }
returnresp, nil
}

var_BlobClient=&localClient{}

// localClient executes the local blob service's code
// to Read or Write bulk files on the current node.
typelocalClientstruct {
localStorage*LocalStorage
}

// NewLocalClient instantiates a local blob service client.
funcNewLocalClient(externalIODirstring) (BlobClient, error) {
storage, err:=NewLocalStorage(externalIODir)
iferr!=nil {
returnnil, errors.Wrap(err, "creating local client")
 }
return&localClient{localStorage: storage}, nil
}

func (c*localClient) ReadFile(
ctx context.Context, filestring, offsetint64,
) (ioctx.ReadCloserCtx, int64, error) {
returnc.localStorage.ReadFile(file, offset)
}

func (c*localClient) Writer(ctx context.Context, filestring) (io.WriteCloser, error) {
returnc.localStorage.Writer(ctx, file)
}

func (c*localClient) List(ctx context.Context, patternstring) ([]string, error) {
returnc.localStorage.List(pattern)
}

func (c*localClient) Delete(ctx context.Context, filestring) error {
returnc.localStorage.Delete(file)
}

func (c*localClient) Stat(ctx context.Context, filestring) (*blobspb.BlobStat, error) {
returnc.localStorage.Stat(file)
}

// BlobClientFactory creates a blob client based on the nodeID we are dialing.
typeBlobClientFactoryfunc(ctx context.Context, dialTarget roachpb.NodeID) (BlobClient, error)

// NewBlobClientFactory returns a BlobClientFactory.
//
// externalIODIr is used to construct a local client when the local
// node is being targetted and the fast path is allowed.
//
// allowLocalFastpath indicates whether the client should create a
// local client (which doesn't go through the network). The fast path
// skips client capability checking so should be used with care.
funcNewBlobClientFactory(
localNodeIDContainer*base.SQLIDContainer,
dialer*nodedialer.Dialer,
externalIODirstring,
allowLocalFastpathbool,
) BlobClientFactory {
returnfunc(ctx context.Context, dialTarget roachpb.NodeID) (BlobClient, error) {
localNodeID, ok:=localNodeIDContainer.OptionalNodeID()
ifok&&dialTarget==0 {
dialTarget=localNodeID
 } elseifdialTarget==0 {
returnnil, errors.New("node ID 0 not supported")
 }

iflocalNodeID==dialTarget&&allowLocalFastpath {
returnNewLocalClient(externalIODir)
 }
conn, err:=dialer.Dial(ctx, dialTarget, rpc.DefaultClass)
iferr!=nil {
returnnil, errors.Wrapf(err, "connecting to node %d", dialTarget)
 }
returnnewRemoteClient(blobspb.NewBlobClient(conn)), nil
 }
}

// NewLocalOnlyBlobClientFactory returns a BlobClientFactory that only
// handles requests for the local node, identified by NodeID = 0.
funcNewLocalOnlyBlobClientFactory(externalIODirstring) BlobClientFactory {
returnfunc(ctx context.Context, dialing roachpb.NodeID) (BlobClient, error) {
ifdialing==0 {
returnNewLocalClient(externalIODir)
 }
returnnil, errors.Errorf("connecting to remote node not supported")
 }
}