list.go

package client

import (
	"context"
	"errors"
	"sync/atomic"
	"time"

	pb "github.com/LiveRamp/gazette/v2/pkg/protocol"
	log "github.com/sirupsen/logrus"
	"google.golang.org/grpc"
)

// PolledList performs periodic polls of a ListRequest. Its most recent
// polled result may be accessed via List.
type PolledList struct {
	ctx    context.Context
	client pb.JournalClient
	req    pb.ListRequest
	resp   atomic.Value
}

// NewPolledList returns a PolledList of the ListRequest which is initialized and
// ready for immediate use, and which will regularly refresh with interval |dur|.
// An error encountered in the first List RPC is returned. Subsequent RPC errors
// will be logged as warnings and retried as part of regular refreshes.
func NewPolledList(ctx context.Context, client pb.JournalClient, dur time.Duration, req pb.ListRequest) (*PolledList, error) {
	var resp, err = ListAllJournals(ctx, client, req)
	if err != nil {
		return nil, err
	}
	var pl = &PolledList{ctx: ctx, client: client, req: req}
	pl.resp.Store(resp)

	go pl.periodicRefresh(dur)
	return pl, nil
}

// List returns the most recent ListResponse.
func (pl *PolledList) List() *pb.ListResponse { return pl.resp.Load().(*pb.ListResponse) }

func (pl *PolledList) periodicRefresh(dur time.Duration) {
	var ticker = time.NewTicker(dur)
	for {
		select {
		case <-ticker.C:
			var resp, err = ListAllJournals(pl.ctx, pl.client, pl.req)
			if err != nil {
				log.WithFields(log.Fields{"err": err, "req": pl.req.String()}).
					Warn("periodic List refresh failed (will retry)")
			} else {
				pl.resp.Store(resp)
			}
		case <-pl.ctx.Done():
			ticker.Stop()
			return
		}
	}
}

// ListAllJournals performs multiple List RPCs, as required to join across multiple
// ListResponse pages, and returns the complete ListResponse of the ListRequest.
// Any encountered error is returned.
func ListAllJournals(ctx context.Context, client pb.JournalClient, req pb.ListRequest) (*pb.ListResponse, error) {
	var resp *pb.ListResponse

	for {
		// List RPCs may be dispatched to any broker.
		if r, err := client.List(pb.WithDispatchDefault(ctx), &req, grpc.FailFast(false)); err != nil {
			return resp, mapGRPCCtxErr(ctx, err)
		} else if err = r.Validate(); err != nil {
			return resp, err
		} else if r.Status != pb.Status_OK {
			return resp, errors.New(r.Status.String())
		} else {
			req.PageToken, r.NextPageToken = r.NextPageToken, ""

			if resp == nil {
				resp = r
			} else {
				resp.Journals = append(resp.Journals, r.Journals...)
			}
		}
		if req.PageToken == "" {
			break // All done.
		}
	}

	if dr, ok := client.(pb.DispatchRouter); ok {
		for _, j := range resp.Journals {
			dr.UpdateRoute(j.Spec.Name.String(), &j.Route)
		}
	}
	return resp, nil
}

// ApplyJournals invokes the Apply RPC.
func ApplyJournals(ctx context.Context, jc pb.JournalClient, req *pb.ApplyRequest) (*pb.ApplyResponse, error) {
	return ApplyJournalsLimit(ctx, jc, req, 0)
}

// ApplyJournalsLimit is a helper function for applying changes to journals which
// may be larger than the configured etcd transaction size limit. The changes in
// |parentReq| will be sent serially in batches of size |maxTxnSize|. If
// |maxTxnSize| is 0 all changes will be attempted as part of a single
// transaction. This function will return the response of the final
// ShardClient.Apply call. Response validation or !OK status from Apply RPC are
// mapped to error. In the event of an error any unapplied changes will be
// available on |parentReq|.
func ApplyJournalsLimit(
	ctx context.Context,
	jc pb.JournalClient,
	parentReq *pb.ApplyRequest,
	maxTxnSize int,
) (*pb.ApplyResponse, error) {
	var changes []pb.ApplyRequest_Change
	if maxTxnSize == 0 {
		maxTxnSize = len(parentReq.Changes)
	}
	var finalResp *pb.ApplyResponse

	for len(parentReq.Changes) > 0 {
		if len(parentReq.Changes) > maxTxnSize {
			changes = parentReq.Changes[:maxTxnSize]
		} else {
			changes = parentReq.Changes
		}

		var req = &pb.ApplyRequest{}
		for _, change := range changes {
			req.Changes = append(req.Changes, change)
		}

		var resp *pb.ApplyResponse
		var err error
		if resp, err = jc.Apply(pb.WithDispatchDefault(ctx), req, grpc.FailFast(false)); err != nil {
			return resp, err
		} else if err = resp.Validate(); err != nil {
			return resp, err
		} else if resp.Status != pb.Status_OK {
			return resp, errors.New(resp.Status.String())
		}
		finalResp = resp
		parentReq.Changes = parentReq.Changes[len(changes):]
	}
	return finalResp, nil
}

// ListAllFragments performs multiple Fragments RPCs, as required to join across multiple
// FragmentsResponse pages, and returns the completed FragmentResponse.
// Any encountered error is returned.
func ListAllFragments(ctx context.Context, client pb.RoutedJournalClient, req pb.FragmentsRequest) (*pb.FragmentsResponse, error) {
	var resp *pb.FragmentsResponse
	var routedCtx = pb.WithDispatchItemRoute(ctx, client, req.Journal.String(), false)

	for {
		if r, err := client.Fragments(routedCtx, &req); err != nil {
			return resp, mapGRPCCtxErr(ctx, err)
		} else if err = r.Validate(); err != nil {
			return resp, err
		} else if r.Status != pb.Status_OK {
			return resp, errors.New(r.Status.String())
		} else {
			req.NextPageToken, r.NextPageToken = r.NextPageToken, 0

			if resp == nil {
				resp = r
			} else {
				resp.Fragments = append(resp.Fragments, r.Fragments...)
			}
			if req.NextPageToken == 0 {
				break // All done.
			}
		}
	}

	return resp, nil
}