<commit_before>// Copyright 2017 The LUCI Authors.
//
// Licensed 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 buildstore

import (
	"net/http"
	"sort"
	"strconv"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/memcache"
	"go.chromium.org/luci/buildbucket"
	bbapi "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/milo/api/buildbot"
	"go.chromium.org/luci/milo/common"
)

// Ternary has 3 defined values: either (zero), yes and no.
type Ternary int

const (
	Either Ternary = iota
	Yes
	No
)

func (t Ternary) filter(q *datastore.Query, fieldName string) *datastore.Query {
	switch t {
	case Yes:
		return q.Eq(fieldName, true)
	case No:
		return q.Eq(fieldName, false)
	default:
		return q
	}
}

// Query is a build query.
type Query struct {
	Master   string
	Builder  string
	Limit    int
	Finished Ternary
	Cursor   string

	// The following fields are tuning parameters specific to a buildstore
	// implementation. Their usage implies understanding of how emulation
	// works.

	// KeyOnly, if true, makes the datastore query keys-only.
	// Loaded Buildbot builds will have only master, builder and number.
	KeyOnly bool // make the data

	// NoAnnotationFetch, if true, will not fetch annotation proto from LogDog.
	// Loaded LUCI builds will not have properties, steps, logs or text.
	NoAnnotationFetch bool

	// NoChangeFetch, if true, will not load change history from Gitiles.
	// Loaded LUCI builds will not have Blame or SourceStamp.Changes.
	NoChangeFetch bool
}

func (q *Query) dsQuery() *datastore.Query {
	dsq := datastore.NewQuery(buildKind)
	if q.Master != "" {
		dsq = dsq.Eq("master", q.Master)
	}
	if q.Builder != "" {
		dsq = dsq.Eq("builder", q.Builder)
	}
	dsq = q.Finished.filter(dsq, "finished")
	if q.Limit > 0 {
		dsq = dsq.Limit(int32(q.Limit))
	}
	if q.KeyOnly {
		dsq = dsq.KeysOnly(true)
	}
	return dsq
}

// QueryResult is a result of running a Query.
type QueryResult struct {
	Builds     []*buildbot.Build // ordered from greater-number to lower-number
	NextCursor string
	PrevCursor string
}

// GetBuilds executes a build query and returns results.
// Does not check access.
func GetBuilds(c context.Context, q Query) (*QueryResult, error) {
	switch {
	case q.Master == "":
		return nil, errors.New("master is required")
	case q.Builder == "":
		return nil, errors.New("builder is required")
	}

	if !EmulationEnabled(c) {
		return getDatastoreBuilds(c, q, true)
	}

	var emulatedBuilds, buildbotBuilds []*buildbot.Build
	err := parallel.FanOutIn(func(work chan<- func() error) {
		work <- func() (err error) {
			res, err := getDatastoreBuilds(c, q, false)
			if res != nil {
				buildbotBuilds = res.Builds
			}
			return
		}
		work <- func() (err error) {
			emulatedBuilds, err = getEmulatedBuilds(c, q)
			return
		}
	})
	if err != nil {
		return nil, errors.Annotate(err, "could not load builds").Err()
	}

	mergedBuilds := mergeBuilds(emulatedBuilds, buildbotBuilds)
	if q.Limit > 0 && len(mergedBuilds) > q.Limit {
		mergedBuilds = mergedBuilds[:q.Limit]
	}
	return &QueryResult{Builds: mergedBuilds}, nil
}

// mergeBuilds merges builds from a and b to one slice.
// The returned builds are ordered by build numbers, descending.
//
// If a build number is present in both a and b, b's build is ignored.
func mergeBuilds(a, b []*buildbot.Build) []*buildbot.Build {
	ret := make([]*buildbot.Build, len(a), len(a)+len(b))
	copy(ret, a)

	// add builds from b that have unique build numbers.
	aNumbers := make(map[int]struct{}, len(a))
	for _, build := range a {
		aNumbers[build.Number] = struct{}{}
	}
	for _, build := range b {
		if _, ok := aNumbers[build.Number]; !ok {
			ret = append(ret, build)
		}
	}
	sort.Slice(ret, func(i, j int) bool {
		return ret[i].Number > ret[j].Number
	})
	return ret
}

func getEmulatedBuilds(c context.Context, q Query) ([]*buildbot.Build, error) {
	if q.Cursor != "" {
		// build query emulation does not support cursors
		logging.Warningf(c, "ignoring cursor %q", q.Cursor)
		q.Cursor = ""
	}

	bb, err := buildbucketClient(c)
	if err != nil {
		return nil, err
	}

	bucket, err := BucketOf(c, q.Master)
	switch {
	case err != nil:
		return nil, errors.Annotate(err, "could not get bucket of %q", q.Master).Err()
	case bucket == "":
		return nil, nil
	}

	search := bb.Search().
		Bucket(bucket).
		Tag(strpair.Format(buildbucket.TagBuilder, q.Builder)).
		Context(c)
	switch q.Finished {
	case Yes:
		search.Status(bbapi.StatusCompleted)
	case No:
		search.Status(bbapi.StatusFilterIncomplete)
	}

	start := clock.Now(c)
	msgs, err := search.Fetch(q.Limit, nil)
	switch apiErr, _ := err.(*googleapi.Error); {
	case apiErr != nil && apiErr.Code == http.StatusForbidden:
		return nil, errors.Annotate(
			err,
			"%q does not have access to bucket %q",
			auth.CurrentIdentity(c),
			bucket).Tag(common.CodeNoAccess).Err()
	case err != nil:
		return nil, errors.Annotate(err, "searching on buildbucket").Err()
	}

	logging.Infof(c, "buildbucket search took %s", clock.Since(c, start))

	builds := make([]*buildbot.Build, len(msgs))
	start = clock.Now(c)
	err = parallel.WorkPool(10, func(work chan<- func() error) {
		for i, msg := range msgs {
			i := i
			msg := msg
			work <- func() error {
				var buildbucketBuild buildbucket.Build
				if err := buildbucketBuild.ParseMessage(msg); err != nil {
					return err
				}
				// may load annotations from logdog, that's why parallelized.
				b, err := buildFromBuildbucket(c, q.Master, &buildbucketBuild, !q.NoAnnotationFetch)
				if err != nil {
					return err
				}
				builds[i] = b
				return nil
			}
		}
	})
	if err != nil {
		return nil, err
	}
	logging.Infof(c, "conversion from buildbucket builds took %s", clock.Since(c, start))

	if !q.NoChangeFetch && len(builds) > 0 {
		start = clock.Now(c)
		// We need to compute blamelist for multiple builds.
		// 1) We don't have a guarantee that the numbers are contiguous
		// 2) For some builds, we may have cached changes
		// => compute blamelist for each build individually

		// cache build revisions before fetching changes
		// in case build numbers are contiguous.
		caches := make([]memcache.Item, len(builds))
		for i, b := range builds {
			caches[i] = buildRevCache(c, b)
		}
		memcache.Set(c, caches...)

		// compute blamelist serially so that git cache is reused.
		for _, b := range builds {
			if err := blame(c, b); err != nil {
				return nil, errors.Annotate(err, "blamelist computation for build #%d failed", b.Number).Err()
			}
		}

		logging.Infof(c, "blamelist computation took %s", clock.Since(c, start))
	}
	return builds, nil
}

func getDatastoreBuilds(c context.Context, q Query, includeExperimental bool) (*QueryResult, error) {
	var builds []*buildEntity
	if q.Limit > 0 {
		builds = make([]*buildEntity, 0, q.Limit)
	}

	dsq := q.dsQuery()

	if !includeExperimental {
		dsq = dsq.Eq("is_experimental", false)
	}

	// CUSTOM CURSOR.
	// This function uses a custom cursor based on build numbers.
	// A cursor is a build number that defines a page boundary.
	// If >=0, it is the inclusive lower boundary.
	//   Example: cursor="10", means return builds ...12, 11, 10.
	// If <0, it is the exclusive upper boundary, negated.
	//   Example: -10, means return builds 9, 8, 7...
	cursorNumber := 0
	order := "-number"
	reverse := false
	hasCursor := false
	if q.Cursor != "" {
		var err error
		if cursorNumber, err = strconv.Atoi(q.Cursor); err == nil {
			hasCursor = true
			if cursorNumber >= 0 {
				dsq = dsq.Gte("number", cursorNumber)
				order = "number"
				reverse = true
			} else {
				dsq = dsq.Lt("number", -cursorNumber)
			}
		} else {
			return nil, errors.New("bad cursor")
		}
	}
	dsq = dsq.Order(order)

	logging.Debugf(c, "running datastore query: %s", dsq)
	err := datastore.GetAll(c, dsq, &builds)
	if err != nil {
		return nil, err
	}
	if reverse {
		for i, j := 0, len(builds)-1; i < j; i, j = i+1, j-1 {
			builds[i], builds[j] = builds[j], builds[i]
		}
	}
	res := &QueryResult{
		Builds: make([]*buildbot.Build, len(builds)),
	}
	for i, b := range builds {
		res.Builds[i] = (*buildbot.Build)(b)
	}

	// Compute prev and next cursors.
	switch {
	case len(res.Builds) > 0:
		// res.Builds are ordered by numbers descending.

		// previous page must display builds with higher numbers.
		if !hasCursor {
			// do not generate a prev cursor for a non-cursor query
		} else {
			// positive cursors are inclusive
			res.PrevCursor = strconv.Itoa(res.Builds[0].Number + 1)
		}

		// next page must display builds with lower numbers.

		if lastNum := res.Builds[len(res.Builds)-1].Number; lastNum == 0 {
			// this is the first ever build, 0, do not generate a cursor
		} else {
			// negative cursors are exclusive.
			res.NextCursor = strconv.Itoa(-lastNum)
		}

	case cursorNumber > 0:
		// no builds and cursor is the inclusive lower boundary
		// e.g. cursor asks for builds after 10,
		// but there are only 0..5 builds.
		// Make the next cursor for builds <10.
		res.NextCursor = strconv.Itoa(-cursorNumber)

	default:
		// there can't be any builds.
	}

	return res, nil
}
<commit_msg>[milo] log and swallow 403 from buildbucket<commit_after>// Copyright 2017 The LUCI Authors.
//
// Licensed 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 buildstore

import (
	"net/http"
	"sort"
	"strconv"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/memcache"
	"go.chromium.org/luci/buildbucket"
	bbapi "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/milo/api/buildbot"
)

// Ternary has 3 defined values: either (zero), yes and no.
type Ternary int

const (
	Either Ternary = iota
	Yes
	No
)

func (t Ternary) filter(q *datastore.Query, fieldName string) *datastore.Query {
	switch t {
	case Yes:
		return q.Eq(fieldName, true)
	case No:
		return q.Eq(fieldName, false)
	default:
		return q
	}
}

// Query is a build query.
type Query struct {
	Master   string
	Builder  string
	Limit    int
	Finished Ternary
	Cursor   string

	// The following fields are tuning parameters specific to a buildstore
	// implementation. Their usage implies understanding of how emulation
	// works.

	// KeyOnly, if true, makes the datastore query keys-only.
	// Loaded Buildbot builds will have only master, builder and number.
	KeyOnly bool // make the data

	// NoAnnotationFetch, if true, will not fetch annotation proto from LogDog.
	// Loaded LUCI builds will not have properties, steps, logs or text.
	NoAnnotationFetch bool

	// NoChangeFetch, if true, will not load change history from Gitiles.
	// Loaded LUCI builds will not have Blame or SourceStamp.Changes.
	NoChangeFetch bool
}

func (q *Query) dsQuery() *datastore.Query {
	dsq := datastore.NewQuery(buildKind)
	if q.Master != "" {
		dsq = dsq.Eq("master", q.Master)
	}
	if q.Builder != "" {
		dsq = dsq.Eq("builder", q.Builder)
	}
	dsq = q.Finished.filter(dsq, "finished")
	if q.Limit > 0 {
		dsq = dsq.Limit(int32(q.Limit))
	}
	if q.KeyOnly {
		dsq = dsq.KeysOnly(true)
	}
	return dsq
}

// QueryResult is a result of running a Query.
type QueryResult struct {
	Builds     []*buildbot.Build // ordered from greater-number to lower-number
	NextCursor string
	PrevCursor string
}

// GetBuilds executes a build query and returns results.
// Does not check access.
func GetBuilds(c context.Context, q Query) (*QueryResult, error) {
	switch {
	case q.Master == "":
		return nil, errors.New("master is required")
	case q.Builder == "":
		return nil, errors.New("builder is required")
	}

	if !EmulationEnabled(c) {
		return getDatastoreBuilds(c, q, true)
	}

	var emulatedBuilds, buildbotBuilds []*buildbot.Build
	err := parallel.FanOutIn(func(work chan<- func() error) {
		work <- func() (err error) {
			res, err := getDatastoreBuilds(c, q, false)
			if res != nil {
				buildbotBuilds = res.Builds
			}
			return
		}
		work <- func() (err error) {
			emulatedBuilds, err = getEmulatedBuilds(c, q)
			return
		}
	})
	if err != nil {
		return nil, errors.Annotate(err, "could not load builds").Err()
	}

	mergedBuilds := mergeBuilds(emulatedBuilds, buildbotBuilds)
	if q.Limit > 0 && len(mergedBuilds) > q.Limit {
		mergedBuilds = mergedBuilds[:q.Limit]
	}
	return &QueryResult{Builds: mergedBuilds}, nil
}

// mergeBuilds merges builds from a and b to one slice.
// The returned builds are ordered by build numbers, descending.
//
// If a build number is present in both a and b, b's build is ignored.
func mergeBuilds(a, b []*buildbot.Build) []*buildbot.Build {
	ret := make([]*buildbot.Build, len(a), len(a)+len(b))
	copy(ret, a)

	// add builds from b that have unique build numbers.
	aNumbers := make(map[int]struct{}, len(a))
	for _, build := range a {
		aNumbers[build.Number] = struct{}{}
	}
	for _, build := range b {
		if _, ok := aNumbers[build.Number]; !ok {
			ret = append(ret, build)
		}
	}
	sort.Slice(ret, func(i, j int) bool {
		return ret[i].Number > ret[j].Number
	})
	return ret
}

func getEmulatedBuilds(c context.Context, q Query) ([]*buildbot.Build, error) {
	if q.Cursor != "" {
		// build query emulation does not support cursors
		logging.Warningf(c, "ignoring cursor %q", q.Cursor)
		q.Cursor = ""
	}

	bb, err := buildbucketClient(c)
	if err != nil {
		return nil, err
	}

	bucket, err := BucketOf(c, q.Master)
	switch {
	case err != nil:
		return nil, errors.Annotate(err, "could not get bucket of %q", q.Master).Err()
	case bucket == "":
		return nil, nil
	}

	search := bb.Search().
		Bucket(bucket).
		Tag(strpair.Format(buildbucket.TagBuilder, q.Builder)).
		Context(c)
	switch q.Finished {
	case Yes:
		search.Status(bbapi.StatusCompleted)
	case No:
		search.Status(bbapi.StatusFilterIncomplete)
	}

	start := clock.Now(c)
	msgs, err := search.Fetch(q.Limit, nil)
	switch apiErr, _ := err.(*googleapi.Error); {
	case apiErr != nil && apiErr.Code == http.StatusForbidden:
		logging.Warningf(c, "%q does not have access to bucket %q. Returning 0 builds.",
			auth.CurrentIdentity(c),
			bucket)
		return nil, nil
	case err != nil:
		return nil, errors.Annotate(err, "searching on buildbucket").Err()
	}

	logging.Infof(c, "buildbucket search took %s", clock.Since(c, start))

	builds := make([]*buildbot.Build, len(msgs))
	start = clock.Now(c)
	err = parallel.WorkPool(10, func(work chan<- func() error) {
		for i, msg := range msgs {
			i := i
			msg := msg
			work <- func() error {
				var buildbucketBuild buildbucket.Build
				if err := buildbucketBuild.ParseMessage(msg); err != nil {
					return err
				}
				// may load annotations from logdog, that's why parallelized.
				b, err := buildFromBuildbucket(c, q.Master, &buildbucketBuild, !q.NoAnnotationFetch)
				if err != nil {
					return err
				}
				builds[i] = b
				return nil
			}
		}
	})
	if err != nil {
		return nil, err
	}
	logging.Infof(c, "conversion from buildbucket builds took %s", clock.Since(c, start))

	if !q.NoChangeFetch && len(builds) > 0 {
		start = clock.Now(c)
		// We need to compute blamelist for multiple builds.
		// 1) We don't have a guarantee that the numbers are contiguous
		// 2) For some builds, we may have cached changes
		// => compute blamelist for each build individually

		// cache build revisions before fetching changes
		// in case build numbers are contiguous.
		caches := make([]memcache.Item, len(builds))
		for i, b := range builds {
			caches[i] = buildRevCache(c, b)
		}
		memcache.Set(c, caches...)

		// compute blamelist serially so that git cache is reused.
		for _, b := range builds {
			if err := blame(c, b); err != nil {
				return nil, errors.Annotate(err, "blamelist computation for build #%d failed", b.Number).Err()
			}
		}

		logging.Infof(c, "blamelist computation took %s", clock.Since(c, start))
	}
	return builds, nil
}

func getDatastoreBuilds(c context.Context, q Query, includeExperimental bool) (*QueryResult, error) {
	var builds []*buildEntity
	if q.Limit > 0 {
		builds = make([]*buildEntity, 0, q.Limit)
	}

	dsq := q.dsQuery()

	if !includeExperimental {
		dsq = dsq.Eq("is_experimental", false)
	}

	// CUSTOM CURSOR.
	// This function uses a custom cursor based on build numbers.
	// A cursor is a build number that defines a page boundary.
	// If >=0, it is the inclusive lower boundary.
	//   Example: cursor="10", means return builds ...12, 11, 10.
	// If <0, it is the exclusive upper boundary, negated.
	//   Example: -10, means return builds 9, 8, 7...
	cursorNumber := 0
	order := "-number"
	reverse := false
	hasCursor := false
	if q.Cursor != "" {
		var err error
		if cursorNumber, err = strconv.Atoi(q.Cursor); err == nil {
			hasCursor = true
			if cursorNumber >= 0 {
				dsq = dsq.Gte("number", cursorNumber)
				order = "number"
				reverse = true
			} else {
				dsq = dsq.Lt("number", -cursorNumber)
			}
		} else {
			return nil, errors.New("bad cursor")
		}
	}
	dsq = dsq.Order(order)

	logging.Debugf(c, "running datastore query: %s", dsq)
	err := datastore.GetAll(c, dsq, &builds)
	if err != nil {
		return nil, err
	}
	if reverse {
		for i, j := 0, len(builds)-1; i < j; i, j = i+1, j-1 {
			builds[i], builds[j] = builds[j], builds[i]
		}
	}
	res := &QueryResult{
		Builds: make([]*buildbot.Build, len(builds)),
	}
	for i, b := range builds {
		res.Builds[i] = (*buildbot.Build)(b)
	}

	// Compute prev and next cursors.
	switch {
	case len(res.Builds) > 0:
		// res.Builds are ordered by numbers descending.

		// previous page must display builds with higher numbers.
		if !hasCursor {
			// do not generate a prev cursor for a non-cursor query
		} else {
			// positive cursors are inclusive
			res.PrevCursor = strconv.Itoa(res.Builds[0].Number + 1)
		}

		// next page must display builds with lower numbers.

		if lastNum := res.Builds[len(res.Builds)-1].Number; lastNum == 0 {
			// this is the first ever build, 0, do not generate a cursor
		} else {
			// negative cursors are exclusive.
			res.NextCursor = strconv.Itoa(-lastNum)
		}

	case cursorNumber > 0:
		// no builds and cursor is the inclusive lower boundary
		// e.g. cursor asks for builds after 10,
		// but there are only 0..5 builds.
		// Make the next cursor for builds <10.
		res.NextCursor = strconv.Itoa(-cursorNumber)

	default:
		// there can't be any builds.
	}

	return res, nil
}
<|endoftext|>


<commit_before>package mongodb

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/compose/transporter/client"
	"github.com/compose/transporter/commitlog"
	"github.com/compose/transporter/log"
	"github.com/compose/transporter/message"
	"github.com/compose/transporter/message/data"
	"github.com/compose/transporter/message/ops"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	_ client.Reader = &Reader{}

	// DefaultCollectionFilter is an empty map of empty maps
	DefaultCollectionFilter = map[string]CollectionFilter{}
)

// CollectionFilter is just a typed map of strings of map[string]interface{}
type CollectionFilter map[string]interface{}

// Reader implements the behavior defined by client.Reader for interfacing with MongoDB.
type Reader struct {
	tail              bool
	collectionFilters map[string]CollectionFilter
	oplogTimeout      time.Duration
}

func newReader(tail bool, filters map[string]CollectionFilter) client.Reader {
	return &Reader{tail, filters, 5 * time.Second}
}

type resultDoc struct {
	doc bson.M
	c   string
}

type iterationComplete struct {
	oplogTime bson.MongoTimestamp
	c         string
}

func (r *Reader) Read(resumeMap map[string]client.MessageSet, filterFn client.NsFilterFunc) client.MessageChanFunc {
	return func(s client.Session, done chan struct{}) (chan client.MessageSet, error) {
		out := make(chan client.MessageSet)
		session := s.(*Session).mgoSession.Copy()
		go func() {
			defer func() {
				session.Close()
				close(out)
			}()
			log.With("db", session.DB("").Name).Infoln("starting Read func")
			collections, err := r.listCollections(session.Copy(), filterFn)
			if err != nil {
				log.With("db", session.DB("").Name).Errorf("unable to list collections, %s", err)
				return
			}
			var wg sync.WaitGroup
			for _, c := range collections {
				var lastID interface{}
				oplogTime := timeAsMongoTimestamp(time.Now())
				var mode commitlog.Mode // default to Copy
				if m, ok := resumeMap[c]; ok {
					lastID = m.Msg.Data().Get("_id")
					mode = m.Mode
					oplogTime = timeAsMongoTimestamp(time.Unix(m.Timestamp, 0))
				}
				if mode == commitlog.Copy {
					if err := r.iterateCollection(r.iterate(lastID, session.Copy(), c), out, done, int64(oplogTime)>>32); err != nil {
						log.With("db", session.DB("").Name).Errorln(err)
						return
					}
					log.With("db", session.DB("").Name).With("collection", c).Infoln("iterating complete")
				}
				if r.tail {
					wg.Add(1)
					log.With("collection", c).Infof("oplog start timestamp: %d", oplogTime)
					go func(wg *sync.WaitGroup, c string, o bson.MongoTimestamp) {
						defer wg.Done()
						errc := r.tailCollection(c, session.Copy(), o, out, done)
						for err := range errc {
							log.With("db", session.DB("").Name).With("collection", c).Errorln(err)
							return
						}
					}(&wg, c, oplogTime)
				}
			}
			log.With("db", session.DB("").Name).Infoln("Read completed")
			// this will block if we're tailing
			wg.Wait()
			return
		}()

		return out, nil
	}
}

func (r *Reader) listCollections(mgoSession *mgo.Session, filterFn func(name string) bool) ([]string, error) {
	defer mgoSession.Close()
	var colls []string
	db := mgoSession.DB("")
	collections, err := db.CollectionNames()
	if err != nil {
		return colls, err
	}
	log.With("db", db.Name).With("num_collections", len(collections)).Infoln("collection count")
	for _, c := range collections {
		if filterFn(c) && !strings.HasPrefix(c, "system.") {
			log.With("db", db.Name).With("collection", c).Infoln("adding for iteration...")
			colls = append(colls, c)
		} else {
			log.With("db", db.Name).With("collection", c).Infoln("skipping iteration...")
		}
	}
	log.With("db", db.Name).Infoln("done iterating collections")
	return colls, nil
}

func (r *Reader) iterateCollection(in <-chan message.Msg, out chan<- client.MessageSet, done chan struct{}, origOplogTime int64) error {
	for {
		select {
		case msg, ok := <-in:
			if !ok {
				return nil
			}
			out <- client.MessageSet{
				Msg:       msg,
				Timestamp: origOplogTime,
			}
		case <-done:
			return errors.New("iteration cancelled")
		}
	}
}

func (r *Reader) iterate(lastID interface{}, s *mgo.Session, c string) <-chan message.Msg {
	msgChan := make(chan message.Msg)
	go func() {
		defer func() {
			s.Close()
			close(msgChan)
		}()
		db := s.DB("").Name
		canReissueQuery := r.requeryable(c, s)
		for {
			log.With("collection", c).Infoln("iterating...")
			session := s.Copy()
			iter := r.catQuery(c, lastID, session).Iter()
			var result bson.M
			for iter.Next(&result) {
				if id, ok := result["_id"]; ok {
					lastID = id
				}
				msgChan <- message.From(ops.Insert, c, data.Data(result))
				result = bson.M{}
			}
			if err := iter.Err(); err != nil {
				log.With("database", db).With("collection", c).Errorf("error reading, %s", err)
				session.Close()
				if canReissueQuery {
					log.With("database", db).With("collection", c).Errorln("attempting to reissue query")
					time.Sleep(5 * time.Second)
					continue
				}
				return
			}
			iter.Close()
			session.Close()
			return
		}
	}()
	return msgChan
}

func (r *Reader) catQuery(c string, lastID interface{}, mgoSession *mgo.Session) *mgo.Query {
	query := bson.M{}
	if f, ok := r.collectionFilters[c]; ok {
		query = bson.M(f)
	}
	if lastID != nil {
		query["_id"] = bson.M{"$gt": lastID}
	}
	return mgoSession.DB("").C(c).Find(query).Sort("_id")
}

func (r *Reader) requeryable(c string, mgoSession *mgo.Session) bool {
	db := mgoSession.DB("")
	indexes, err := db.C(c).Indexes()
	if err != nil {
		log.With("database", db.Name).With("collection", c).Errorf("unable to list indexes, %s", err)
		return false
	}
	for _, index := range indexes {
		if index.Key[0] == "_id" {
			var result bson.M
			err := db.C(c).Find(nil).Select(bson.M{"_id": 1}).One(&result)
			if err != nil {
				log.With("database", db.Name).With("collection", c).Errorf("unable to sample document, %s", err)
				break
			}
			if id, ok := result["_id"]; ok && sortable(id) {
				return true
			}
			break
		}
	}
	log.With("database", db.Name).With("collection", c).Infoln("invalid _id, any issues copying will be aborted")
	return false
}

func sortable(id interface{}) bool {
	switch id.(type) {
	case bson.ObjectId, string, float64, int64, time.Time:
		return true
	}
	return false
}

func (r *Reader) tailCollection(c string, mgoSession *mgo.Session, oplogTime bson.MongoTimestamp, out chan<- client.MessageSet, done chan struct{}) chan error {
	errc := make(chan error)
	go func() {
		defer func() {
			mgoSession.Close()
			close(errc)
		}()

		var (
			collection = mgoSession.DB("local").C("oplog.rs")
			result     oplogDoc // hold the document
			db         = mgoSession.DB("").Name
			query      = bson.M{"ns": fmt.Sprintf("%s.%s", db, c), "ts": bson.M{"$gte": oplogTime}}
			iter       = collection.Find(query).LogReplay().Sort("$natural").Tail(r.oplogTimeout)
		)
		defer iter.Close()

		for {
			log.With("db", db).Infof("tailing oplog with query %+v", query)
			select {
			case <-done:
				log.With("db", db).Infoln("tailing stopping...")
				return
			default:
				for iter.Next(&result) {
					if result.validOp() {
						var (
							doc bson.M
							err error
							op  ops.Op
						)
						switch result.Op {
						case "i":
							op = ops.Insert
							doc = result.O
						case "d":
							op = ops.Delete
							doc = result.O
						case "u":
							op = ops.Update
							doc, err = r.getOriginalDoc(result.O2, c, mgoSession)
							if err != nil {
								// errors aren't fatal here, but we need to send it down the pipe
								log.With("ns", result.Ns).Errorf("unable to getOriginalDoc, %s", err)
								continue
							}
						}

						msg := message.From(op, c, data.Data(doc)).(*message.Base)
						msg.TS = int64(result.Ts) >> 32

						out <- client.MessageSet{
							Msg:       msg,
							Timestamp: msg.TS,
							Mode:      commitlog.Sync,
						}
						oplogTime = result.Ts
					}
					result = oplogDoc{}
				}
			}

			if iter.Timeout() {
				continue
			}
			if iter.Err() != nil {
				log.With("path", db).Errorf("error tailing oplog, %s", iter.Err())
				// return adaptor.NewError(adaptor.CRITICAL, m.path, fmt.Sprintf("MongoDB error (error reading collection %s)", iter.Err()), nil)
			}

			query = bson.M{"ts": bson.M{"$gte": oplogTime}}
			iter = collection.Find(query).LogReplay().Tail(r.oplogTimeout)
			time.Sleep(100 * time.Millisecond)
		}

	}()
	return errc
}

// getOriginalDoc retrieves the original document from the database.
// transporter has no knowledge of update operations, all updates work as wholesale document replaces
func (r *Reader) getOriginalDoc(doc bson.M, c string, s *mgo.Session) (result bson.M, err error) {
	id, exists := doc["_id"]
	if !exists {
		return result, fmt.Errorf("can't get _id from document")
	}

	query := bson.M{}
	if f, ok := r.collectionFilters[c]; ok {
		query = bson.M(f)
	}
	query["_id"] = id

	err = s.DB("").C(c).Find(query).One(&result)
	if err != nil {
		err = fmt.Errorf("%s.%s %v %v", s.DB("").Name, c, id, err)
	}
	return
}

// oplogDoc are representations of the mongodb oplog document
// detailed here, among other places.  http://www.kchodorow.com/blog/2010/10/12/replication-internals/
type oplogDoc struct {
	Ts bson.MongoTimestamp `bson:"ts"`
	H  int64               `bson:"h"`
	V  int                 `bson:"v"`
	Op string              `bson:"op"`
	Ns string              `bson:"ns"`
	O  bson.M              `bson:"o"`
	O2 bson.M              `bson:"o2"`
}

// validOp checks to see if we're an insert, delete, or update, otherwise the
// document is skilled.
// TODO: skip system collections
func (o *oplogDoc) validOp() bool {
	return o.Op == "i" || o.Op == "d" || o.Op == "u"
}

func timeAsMongoTimestamp(t time.Time) bson.MongoTimestamp {
	return bson.MongoTimestamp(t.Unix() << 32)
}
<commit_msg>also verify oplog doc is for the tailing collection (#386)<commit_after>package mongodb

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/compose/transporter/client"
	"github.com/compose/transporter/commitlog"
	"github.com/compose/transporter/log"
	"github.com/compose/transporter/message"
	"github.com/compose/transporter/message/data"
	"github.com/compose/transporter/message/ops"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	_ client.Reader = &Reader{}

	// DefaultCollectionFilter is an empty map of empty maps
	DefaultCollectionFilter = map[string]CollectionFilter{}
)

// CollectionFilter is just a typed map of strings of map[string]interface{}
type CollectionFilter map[string]interface{}

// Reader implements the behavior defined by client.Reader for interfacing with MongoDB.
type Reader struct {
	tail              bool
	collectionFilters map[string]CollectionFilter
	oplogTimeout      time.Duration
}

func newReader(tail bool, filters map[string]CollectionFilter) client.Reader {
	return &Reader{tail, filters, 5 * time.Second}
}

type resultDoc struct {
	doc bson.M
	c   string
}

type iterationComplete struct {
	oplogTime bson.MongoTimestamp
	c         string
}

func (r *Reader) Read(resumeMap map[string]client.MessageSet, filterFn client.NsFilterFunc) client.MessageChanFunc {
	return func(s client.Session, done chan struct{}) (chan client.MessageSet, error) {
		out := make(chan client.MessageSet)
		session := s.(*Session).mgoSession.Copy()
		go func() {
			defer func() {
				session.Close()
				close(out)
			}()
			log.With("db", session.DB("").Name).Infoln("starting Read func")
			collections, err := r.listCollections(session.Copy(), filterFn)
			if err != nil {
				log.With("db", session.DB("").Name).Errorf("unable to list collections, %s", err)
				return
			}
			var wg sync.WaitGroup
			for _, c := range collections {
				var lastID interface{}
				oplogTime := timeAsMongoTimestamp(time.Now())
				var mode commitlog.Mode // default to Copy
				if m, ok := resumeMap[c]; ok {
					lastID = m.Msg.Data().Get("_id")
					mode = m.Mode
					oplogTime = timeAsMongoTimestamp(time.Unix(m.Timestamp, 0))
				}
				if mode == commitlog.Copy {
					if err := r.iterateCollection(r.iterate(lastID, session.Copy(), c), out, done, int64(oplogTime)>>32); err != nil {
						log.With("db", session.DB("").Name).Errorln(err)
						return
					}
					log.With("db", session.DB("").Name).With("collection", c).Infoln("iterating complete")
				}
				if r.tail {
					wg.Add(1)
					log.With("collection", c).Infof("oplog start timestamp: %d", oplogTime)
					go func(wg *sync.WaitGroup, c string, o bson.MongoTimestamp) {
						defer wg.Done()
						errc := r.tailCollection(c, session.Copy(), o, out, done)
						for err := range errc {
							log.With("db", session.DB("").Name).With("collection", c).Errorln(err)
							return
						}
					}(&wg, c, oplogTime)
				}
			}
			log.With("db", session.DB("").Name).Infoln("Read completed")
			// this will block if we're tailing
			wg.Wait()
			return
		}()

		return out, nil
	}
}

func (r *Reader) listCollections(mgoSession *mgo.Session, filterFn func(name string) bool) ([]string, error) {
	defer mgoSession.Close()
	var colls []string
	db := mgoSession.DB("")
	collections, err := db.CollectionNames()
	if err != nil {
		return colls, err
	}
	log.With("db", db.Name).With("num_collections", len(collections)).Infoln("collection count")
	for _, c := range collections {
		if filterFn(c) && !strings.HasPrefix(c, "system.") {
			log.With("db", db.Name).With("collection", c).Infoln("adding for iteration...")
			colls = append(colls, c)
		} else {
			log.With("db", db.Name).With("collection", c).Infoln("skipping iteration...")
		}
	}
	log.With("db", db.Name).Infoln("done iterating collections")
	return colls, nil
}

func (r *Reader) iterateCollection(in <-chan message.Msg, out chan<- client.MessageSet, done chan struct{}, origOplogTime int64) error {
	for {
		select {
		case msg, ok := <-in:
			if !ok {
				return nil
			}
			out <- client.MessageSet{
				Msg:       msg,
				Timestamp: origOplogTime,
			}
		case <-done:
			return errors.New("iteration cancelled")
		}
	}
}

func (r *Reader) iterate(lastID interface{}, s *mgo.Session, c string) <-chan message.Msg {
	msgChan := make(chan message.Msg)
	go func() {
		defer func() {
			s.Close()
			close(msgChan)
		}()
		db := s.DB("").Name
		canReissueQuery := r.requeryable(c, s)
		for {
			log.With("collection", c).Infoln("iterating...")
			session := s.Copy()
			iter := r.catQuery(c, lastID, session).Iter()
			var result bson.M
			for iter.Next(&result) {
				if id, ok := result["_id"]; ok {
					lastID = id
				}
				msgChan <- message.From(ops.Insert, c, data.Data(result))
				result = bson.M{}
			}
			if err := iter.Err(); err != nil {
				log.With("database", db).With("collection", c).Errorf("error reading, %s", err)
				session.Close()
				if canReissueQuery {
					log.With("database", db).With("collection", c).Errorln("attempting to reissue query")
					time.Sleep(5 * time.Second)
					continue
				}
				return
			}
			iter.Close()
			session.Close()
			return
		}
	}()
	return msgChan
}

func (r *Reader) catQuery(c string, lastID interface{}, mgoSession *mgo.Session) *mgo.Query {
	query := bson.M{}
	if f, ok := r.collectionFilters[c]; ok {
		query = bson.M(f)
	}
	if lastID != nil {
		query["_id"] = bson.M{"$gt": lastID}
	}
	return mgoSession.DB("").C(c).Find(query).Sort("_id")
}

func (r *Reader) requeryable(c string, mgoSession *mgo.Session) bool {
	db := mgoSession.DB("")
	indexes, err := db.C(c).Indexes()
	if err != nil {
		log.With("database", db.Name).With("collection", c).Errorf("unable to list indexes, %s", err)
		return false
	}
	for _, index := range indexes {
		if index.Key[0] == "_id" {
			var result bson.M
			err := db.C(c).Find(nil).Select(bson.M{"_id": 1}).One(&result)
			if err != nil {
				log.With("database", db.Name).With("collection", c).Errorf("unable to sample document, %s", err)
				break
			}
			if id, ok := result["_id"]; ok && sortable(id) {
				return true
			}
			break
		}
	}
	log.With("database", db.Name).With("collection", c).Infoln("invalid _id, any issues copying will be aborted")
	return false
}

func sortable(id interface{}) bool {
	switch id.(type) {
	case bson.ObjectId, string, float64, int64, time.Time:
		return true
	}
	return false
}

func (r *Reader) tailCollection(c string, mgoSession *mgo.Session, oplogTime bson.MongoTimestamp, out chan<- client.MessageSet, done chan struct{}) chan error {
	errc := make(chan error)
	go func() {
		defer func() {
			mgoSession.Close()
			close(errc)
		}()

		var (
			collection = mgoSession.DB("local").C("oplog.rs")
			result     oplogDoc // hold the document
			db         = mgoSession.DB("").Name
			ns         = fmt.Sprintf("%s.%s", db, c)
			query      = bson.M{"ns": ns, "ts": bson.M{"$gte": oplogTime}}
			iter       = collection.Find(query).LogReplay().Sort("$natural").Tail(r.oplogTimeout)
		)
		defer iter.Close()

		for {
			log.With("db", db).Infof("tailing oplog with query %+v", query)
			select {
			case <-done:
				log.With("db", db).Infoln("tailing stopping...")
				return
			default:
				for iter.Next(&result) {
					if result.validOp(ns) {
						var (
							doc bson.M
							err error
							op  ops.Op
						)
						switch result.Op {
						case "i":
							op = ops.Insert
							doc = result.O
						case "d":
							op = ops.Delete
							doc = result.O
						case "u":
							op = ops.Update
							doc, err = r.getOriginalDoc(result.O2, c, mgoSession)
							if err != nil {
								// errors aren't fatal here, but we need to send it down the pipe
								log.With("ns", result.Ns).Errorf("unable to getOriginalDoc, %s", err)
								continue
							}
						}

						msg := message.From(op, c, data.Data(doc)).(*message.Base)
						msg.TS = int64(result.Ts) >> 32

						out <- client.MessageSet{
							Msg:       msg,
							Timestamp: msg.TS,
							Mode:      commitlog.Sync,
						}
						oplogTime = result.Ts
					}
					result = oplogDoc{}
				}
			}

			if iter.Timeout() {
				continue
			}
			if iter.Err() != nil {
				log.With("path", db).Errorf("error tailing oplog, %s", iter.Err())
				// return adaptor.NewError(adaptor.CRITICAL, m.path, fmt.Sprintf("MongoDB error (error reading collection %s)", iter.Err()), nil)
			}

			query = bson.M{"ts": bson.M{"$gte": oplogTime}}
			iter = collection.Find(query).LogReplay().Tail(r.oplogTimeout)
			time.Sleep(100 * time.Millisecond)
		}

	}()
	return errc
}

// getOriginalDoc retrieves the original document from the database.
// transporter has no knowledge of update operations, all updates work as wholesale document replaces
func (r *Reader) getOriginalDoc(doc bson.M, c string, s *mgo.Session) (result bson.M, err error) {
	id, exists := doc["_id"]
	if !exists {
		return result, fmt.Errorf("can't get _id from document")
	}

	query := bson.M{}
	if f, ok := r.collectionFilters[c]; ok {
		query = bson.M(f)
	}
	query["_id"] = id

	err = s.DB("").C(c).Find(query).One(&result)
	if err != nil {
		err = fmt.Errorf("%s.%s %v %v", s.DB("").Name, c, id, err)
	}
	return
}

// oplogDoc are representations of the mongodb oplog document
// detailed here, among other places.  http://www.kchodorow.com/blog/2010/10/12/replication-internals/
type oplogDoc struct {
	Ts bson.MongoTimestamp `bson:"ts"`
	H  int64               `bson:"h"`
	V  int                 `bson:"v"`
	Op string              `bson:"op"`
	Ns string              `bson:"ns"`
	O  bson.M              `bson:"o"`
	O2 bson.M              `bson:"o2"`
}

// validOp checks to see if we're an insert, delete, or update, otherwise the
// document is skilled.
func (o *oplogDoc) validOp(ns string) bool {
	return ns == o.Ns && (o.Op == "i" || o.Op == "d" || o.Op == "u")
}

func timeAsMongoTimestamp(t time.Time) bson.MongoTimestamp {
	return bson.MongoTimestamp(t.Unix() << 32)
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package network

import (
	"errors"
	"fmt"
	"math"

	"github.com/ava-labs/avalanchego/utils/wrappers"
)

var (
	errMissingField = errors.New("message missing field")
	errBadOp        = errors.New("input field has invalid operation")
)

// Codec defines the serialization and deserialization of network messages
type Codec struct {
	compressor Compressor
}

// Pack attempts to pack a map of fields into a message.
// The first byte of the message is the opcode of the message.
// Uses [buffer] to hold the message's byte repr.
// [buffer]'s contents may be overwritten.
// [buffer] may be nil.
func (c Codec) Pack(buffer []byte, op Op, fieldValues map[Field]interface{}, compress bool) (Msg, error) {
	msgFields, ok := Messages[op]
	if !ok {
		return nil, errBadOp
	}

	p := wrappers.Packer{
		MaxSize: math.MaxInt32,
		Bytes:   buffer[:0],
	}
	// Pack the op code (message type)
	p.PackByte(byte(op))

	// If messages of this type may be compressed, pack whether the payload is compressed
	compressableType := op != Version && op != GetVersion // TODO in the future, always pack
	if compressableType {
		p.PackBool(compress)
	}

	// Pack the payload
	for _, field := range msgFields {
		data, ok := fieldValues[field]
		if !ok {
			return nil, errMissingField
		}
		field.Packer()(&p, data)
	}
	if p.Err != nil {
		return nil, p.Err
	}
	msg := &msg{
		op:     op,
		fields: fieldValues,
		bytes:  p.Bytes,
	}
	if !compress {
		return msg, nil
	}

	// If [compress], compress the payload (not the op code, not isCompressed)
	payloadBytes := msg.bytes[wrappers.BoolLen+wrappers.ByteLen:]
	compressedPayload, err := c.compressor.Compress(payloadBytes)
	if err != nil {
		return nil, fmt.Errorf("couldn't compress payload of %s message: %s", op, err)
	}
	// Shrink [msg.bytes] to be the right size:
	// 1 byte for the op code, 1 byte for isCompressed, len(compressedPayload) bytes for the compressed payload
	msg.bytes = msg.bytes[:wrappers.BoolLen+wrappers.ByteLen+len(compressedPayload)]
	// Copy the compressed payload into the message bytes
	copy(msg.bytes[wrappers.BoolLen+wrappers.ByteLen:], compressedPayload)
	return msg, nil
}

// Parse attempts to convert bytes into a message.
// The first byte of the message is the opcode of the message.
func (c Codec) Parse(b []byte, mayBeCompressed bool) (Msg, error) {
	p := wrappers.Packer{Bytes: b}

	// Unpack the op code (message type)
	op := Op(p.UnpackByte())

	msgFields, ok := Messages[op]
	if !ok { // Unknown message type
		return nil, errBadOp
	}

	// See if messages of this type may be compressed
	compressableType := op != Version && op != GetVersion
	compressed := false
	if compressableType && mayBeCompressed {
		compressed = p.UnpackBool()
	}
	if p.Err != nil {
		return nil, p.Err
	}

	// If the payload is compressed, decompress it
	if compressed {
		// The slice below is guaranteed to be in-bounds because [p.Err] == nil
		compressedPayloadBytes := p.Bytes[wrappers.ByteLen+wrappers.BoolLen:]
		payloadBytes, err := c.compressor.Decompress(compressedPayloadBytes)
		if err != nil {
			return nil, fmt.Errorf("couldn't decompress payload of %s message: %s", op, err)
		}
		// Replace the compressed payload with the decompressedPayload.
		// Remove the compressed payload.
		p.Bytes = p.Bytes[:wrappers.ByteLen+wrappers.BoolLen]
		// Attach the decompressed payload.
		p.Bytes = append(p.Bytes, payloadBytes...)
	}

	// Parse each field of the payload
	fieldValues := make(map[Field]interface{}, len(msgFields))
	for _, field := range msgFields {
		fieldValues[field] = field.Unpacker()(&p)
	}

	if p.Offset != len(b) {
		p.Add(fmt.Errorf("expected length %d got %d", len(b), p.Offset))
	}

	return &msg{
		op:     op,
		fields: fieldValues,
		bytes:  b,
	}, p.Err
}
<commit_msg>comment<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package network

import (
	"errors"
	"fmt"
	"math"

	"github.com/ava-labs/avalanchego/utils/wrappers"
)

var (
	errMissingField = errors.New("message missing field")
	errBadOp        = errors.New("input field has invalid operation")
)

// Codec defines the serialization and deserialization of network messages
type Codec struct {
	compressor Compressor
}

// Pack attempts to pack a map of fields into a message.
// The first byte of the message is the opcode of the message.
// Uses [buffer] to hold the message's byte repr.
// [buffer]'s contents may be overwritten.
// [buffer] may be nil.
func (c Codec) Pack(buffer []byte, op Op, fieldValues map[Field]interface{}, compress bool) (Msg, error) {
	msgFields, ok := Messages[op]
	if !ok {
		return nil, errBadOp
	}

	p := wrappers.Packer{
		MaxSize: math.MaxInt32,
		Bytes:   buffer[:0],
	}
	// Pack the op code (message type)
	p.PackByte(byte(op))

	// If messages of this type may be compressed, pack whether the payload is compressed
	compressableType := op != Version && op != GetVersion // TODO in the future, always pack
	if compressableType {
		p.PackBool(compress)
	}

	// Pack the payload
	for _, field := range msgFields {
		data, ok := fieldValues[field]
		if !ok {
			return nil, errMissingField
		}
		field.Packer()(&p, data)
	}
	if p.Err != nil {
		return nil, p.Err
	}
	msg := &msg{
		op:     op,
		fields: fieldValues,
		bytes:  p.Bytes,
	}
	if !compress {
		return msg, nil
	}

	// If [compress], compress the payload (not the op code, not isCompressed).
	// The slice below is guaranteed to be in-bounds because [p.Err] == nil
	// implies that len(msg.bytes) >= 2
	payloadBytes := msg.bytes[wrappers.BoolLen+wrappers.ByteLen:]
	compressedPayloadBytes, err := c.compressor.Compress(payloadBytes)
	if err != nil {
		return nil, fmt.Errorf("couldn't compress payload of %s message: %s", op, err)
	}
	// Remove the payload
	msg.bytes = msg.bytes[:wrappers.BoolLen+wrappers.ByteLen]
	// Attach the compressed payload
	msg.bytes = append(msg.bytes, compressedPayloadBytes...)
	return msg, nil
}

// Parse attempts to convert bytes into a message.
// The first byte of the message is the opcode of the message.
func (c Codec) Parse(b []byte, mayBeCompressed bool) (Msg, error) {
	p := wrappers.Packer{Bytes: b}

	// Unpack the op code (message type)
	op := Op(p.UnpackByte())

	msgFields, ok := Messages[op]
	if !ok { // Unknown message type
		return nil, errBadOp
	}

	// See if messages of this type may be compressed
	compressableType := op != Version && op != GetVersion
	compressed := false
	if compressableType && mayBeCompressed {
		compressed = p.UnpackBool()
	}
	if p.Err != nil {
		return nil, p.Err
	}

	// If the payload is compressed, decompress it
	if compressed {
		// The slice below is guaranteed to be in-bounds because [p.Err] == nil
		compressedPayloadBytes := p.Bytes[wrappers.ByteLen+wrappers.BoolLen:]
		payloadBytes, err := c.compressor.Decompress(compressedPayloadBytes)
		if err != nil {
			return nil, fmt.Errorf("couldn't decompress payload of %s message: %s", op, err)
		}
		// Replace the compressed payload with the decompressed payload.
		// Remove the compressed payload.
		p.Bytes = p.Bytes[:wrappers.ByteLen+wrappers.BoolLen]
		// Attach the decompressed payload.
		p.Bytes = append(p.Bytes, payloadBytes...)
	}

	// Parse each field of the payload
	fieldValues := make(map[Field]interface{}, len(msgFields))
	for _, field := range msgFields {
		fieldValues[field] = field.Unpacker()(&p)
	}

	if p.Offset != len(b) {
		p.Add(fmt.Errorf("expected length %d got %d", len(b), p.Offset))
	}

	return &msg{
		op:     op,
		fields: fieldValues,
		bytes:  b,
	}, p.Err
}
<|endoftext|>


<commit_before>package array

func sortArrayByParity(arr []int) []int {
	return useOnePass(arr)
}

// useOnePass time complexity O(N), space complexity O(1)
func useOnePass(arr []int) []int {
	n := len(arr)
	l, r := 0, n-1
	for l < r {
		if arr[l]%2 == 0 {
			l++
		} else {
			arr[l], arr[r] = arr[r], arr[l]
		}
		if arr[r]%2 != 0 {
			r--
		} else {
			arr[r], arr[l] = arr[l], arr[r]
		}
	}
	return arr
}
<commit_msg>remove duplicated code in 905<commit_after>package array

func sortArrayByParity(arr []int) []int {
	return useOnePass(arr)
}

// useOnePass time complexity O(N), space complexity O(1)
func useOnePass(arr []int) []int {
	n := len(arr)
	l, r := 0, n-1
	for l < r {
		if arr[l]%2 == 0 {
			l++
		} else {
			arr[l], arr[r] = arr[r], arr[l]
			// arr[r] is odd now
			r--
		}
	}
	return arr
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"google.golang.org/api/cloudkms/v1"
)

type kmsKeyRingId struct {
	Project  string
	Location string
	Name     string
}

func (s *kmsKeyRingId) keyRingId() string {
	return fmt.Sprintf("projects/%s/locations/%s/keyRings/%s", s.Project, s.Location, s.Name)
}

func (s *kmsKeyRingId) terraformId() string {
	return fmt.Sprintf("%s/%s/%s", s.Project, s.Location, s.Name)
}

func parseKmsKeyRingId(id string, config *Config) (*kmsKeyRingId, error) {
	parts := strings.Split(id, "/")

	keyRingIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	keyRingIdWithoutProjectRegex := regexp.MustCompile("^([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	keyRingRelativeLinkRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-]+)/keyRings/([a-zA-Z0-9_-]{1,63})$")

	if keyRingIdRegex.MatchString(id) {
		return &kmsKeyRingId{
			Project:  parts[0],
			Location: parts[1],
			Name:     parts[2],
		}, nil
	}

	if keyRingIdWithoutProjectRegex.MatchString(id) {
		if config.Project == "" {
			return nil, fmt.Errorf("The default project for the provider must be set when using the `{location}/{keyRingName}` id format.")
		}

		return &kmsKeyRingId{
			Project:  config.Project,
			Location: parts[0],
			Name:     parts[1],
		}, nil
	}

	if parts := keyRingRelativeLinkRegex.FindStringSubmatch(id); parts != nil {
		return &kmsKeyRingId{
			Project:  parts[1],
			Location: parts[2],
			Name:     parts[3],
		}, nil
	}
	return nil, fmt.Errorf("Invalid KeyRing id format, expecting `{projectId}/{locationId}/{keyRingName}` or `{locationId}/{keyRingName}.`")
}

func kmsCryptoKeyRingsEquivalent(k, old, new string, d *schema.ResourceData) bool {
	keyRingIdWithSpecifiersRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-])+/keyRings/([a-zA-Z0-9_-]{1,63})$")
	normalizedKeyRingIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	if matches := keyRingIdWithSpecifiersRegex.FindStringSubmatch(new); matches != nil {
		normMatches := normalizedKeyRingIdRegex.FindStringSubmatch(old)
		return normMatches != nil && normMatches[1] == matches[1] && normMatches[2] == matches[2] && normMatches[3] == matches[3]
	}
	return false
}

type kmsCryptoKeyId struct {
	KeyRingId kmsKeyRingId
	Name      string
}

func (s *kmsCryptoKeyId) cryptoKeyId() string {
	return fmt.Sprintf("%s/cryptoKeys/%s", s.KeyRingId.keyRingId(), s.Name)
}

func (s *kmsCryptoKeyId) terraformId() string {
	return fmt.Sprintf("%s/%s", s.KeyRingId.terraformId(), s.Name)
}

func validateKmsCryptoKeyRotationPeriod(value interface{}, _ string) (ws []string, errors []error) {
	period := value.(string)
	pattern := regexp.MustCompile(`^([0-9.]*\d)s$`)
	match := pattern.FindStringSubmatch(period)

	if len(match) == 0 {
		errors = append(errors, fmt.Errorf("Invalid rotation period format: %s", period))
		// Cannot continue to validate because we cannot extract a number.
		return
	}

	number := match[1]
	seconds, err := strconv.ParseFloat(number, 64)

	if err != nil {
		errors = append(errors, err)
	} else {
		if seconds < 86400.0 {
			errors = append(errors, fmt.Errorf("Rotation period must be greater than one day"))
		}

		parts := strings.Split(number, ".")

		if len(parts) > 1 && len(parts[1]) > 9 {
			errors = append(errors, fmt.Errorf("Rotation period cannot have more than 9 fractional digits"))
		}
	}

	return
}

func kmsCryptoKeyNextRotation(now time.Time, period string) (result string, err error) {
	var duration time.Duration

	duration, err = time.ParseDuration(period)

	if err == nil {
		result = now.UTC().Add(duration).Format(time.RFC3339Nano)
	}

	return
}

func parseKmsCryptoKeyId(id string, config *Config) (*kmsCryptoKeyId, error) {
	parts := strings.Split(id, "/")

	cryptoKeyIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})/([a-zA-Z0-9_-]{1,63})$")
	cryptoKeyIdWithoutProjectRegex := regexp.MustCompile("^([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})/([a-zA-Z0-9_-]{1,63})$")
	cryptoKeyRelativeLinkRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-]+)/keyRings/([a-zA-Z0-9_-]{1,63})/cryptoKeys/([a-zA-Z0-9_-]{1,63})$")

	if cryptoKeyIdRegex.MatchString(id) {
		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  parts[0],
				Location: parts[1],
				Name:     parts[2],
			},
			Name: parts[3],
		}, nil
	}

	if cryptoKeyIdWithoutProjectRegex.MatchString(id) {
		if config.Project == "" {
			return nil, fmt.Errorf("The default project for the provider must be set when using the `{location}/{keyRingName}/{cryptoKeyName}` id format.")
		}

		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  config.Project,
				Location: parts[0],
				Name:     parts[1],
			},
			Name: parts[2],
		}, nil
	}

	if parts := cryptoKeyRelativeLinkRegex.FindStringSubmatch(id); parts != nil {
		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  parts[1],
				Location: parts[2],
				Name:     parts[3],
			},
			Name: parts[4],
		}, nil
	}
	return nil, fmt.Errorf("Invalid CryptoKey id format, expecting `{projectId}/{locationId}/{KeyringName}/{cryptoKeyName}` or `{locationId}/{keyRingName}/{cryptoKeyName}, got id: %s`", id)
}

func clearCryptoKeyVersions(cryptoKeyId *kmsCryptoKeyId, userAgent string, config *Config) error {
	versionsClient := config.NewKmsClient(userAgent).Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions

	listCall := versionsClient.List(cryptoKeyId.cryptoKeyId())
	if config.UserProjectOverride {
		listCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
	}
	versionsResponse, err := listCall.Do()

	if err != nil {
		return err
	}

	for _, version := range versionsResponse.CryptoKeyVersions {
		request := &cloudkms.DestroyCryptoKeyVersionRequest{}
		destroyCall := versionsClient.Destroy(version.Name, request)
		if config.UserProjectOverride {
			destroyCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
		}
		_, err = destroyCall.Do()

		if err != nil {
			return err
		}
	}

	return nil
}

func disableCryptoKeyRotation(cryptoKeyId *kmsCryptoKeyId, userAgent string, config *Config) error {
	keyClient := config.NewKmsClient(userAgent).Projects.Locations.KeyRings.CryptoKeys
	patchCall := keyClient.Patch(cryptoKeyId.cryptoKeyId(), &cloudkms.CryptoKey{
		NullFields: []string{"rotationPeriod", "nextRotationTime"},
	}).
		UpdateMask("rotationPeriod,nextRotationTime")
	if config.UserProjectOverride {
		patchCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
	}
	_, err := patchCall.Do()

	return err
}
<commit_msg>skip destroyed key versions (#6669)<commit_after>package google

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"google.golang.org/api/cloudkms/v1"
)

type kmsKeyRingId struct {
	Project  string
	Location string
	Name     string
}

func (s *kmsKeyRingId) keyRingId() string {
	return fmt.Sprintf("projects/%s/locations/%s/keyRings/%s", s.Project, s.Location, s.Name)
}

func (s *kmsKeyRingId) terraformId() string {
	return fmt.Sprintf("%s/%s/%s", s.Project, s.Location, s.Name)
}

func parseKmsKeyRingId(id string, config *Config) (*kmsKeyRingId, error) {
	parts := strings.Split(id, "/")

	keyRingIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	keyRingIdWithoutProjectRegex := regexp.MustCompile("^([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	keyRingRelativeLinkRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-]+)/keyRings/([a-zA-Z0-9_-]{1,63})$")

	if keyRingIdRegex.MatchString(id) {
		return &kmsKeyRingId{
			Project:  parts[0],
			Location: parts[1],
			Name:     parts[2],
		}, nil
	}

	if keyRingIdWithoutProjectRegex.MatchString(id) {
		if config.Project == "" {
			return nil, fmt.Errorf("The default project for the provider must be set when using the `{location}/{keyRingName}` id format.")
		}

		return &kmsKeyRingId{
			Project:  config.Project,
			Location: parts[0],
			Name:     parts[1],
		}, nil
	}

	if parts := keyRingRelativeLinkRegex.FindStringSubmatch(id); parts != nil {
		return &kmsKeyRingId{
			Project:  parts[1],
			Location: parts[2],
			Name:     parts[3],
		}, nil
	}
	return nil, fmt.Errorf("Invalid KeyRing id format, expecting `{projectId}/{locationId}/{keyRingName}` or `{locationId}/{keyRingName}.`")
}

func kmsCryptoKeyRingsEquivalent(k, old, new string, d *schema.ResourceData) bool {
	keyRingIdWithSpecifiersRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-])+/keyRings/([a-zA-Z0-9_-]{1,63})$")
	normalizedKeyRingIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})$")
	if matches := keyRingIdWithSpecifiersRegex.FindStringSubmatch(new); matches != nil {
		normMatches := normalizedKeyRingIdRegex.FindStringSubmatch(old)
		return normMatches != nil && normMatches[1] == matches[1] && normMatches[2] == matches[2] && normMatches[3] == matches[3]
	}
	return false
}

type kmsCryptoKeyId struct {
	KeyRingId kmsKeyRingId
	Name      string
}

func (s *kmsCryptoKeyId) cryptoKeyId() string {
	return fmt.Sprintf("%s/cryptoKeys/%s", s.KeyRingId.keyRingId(), s.Name)
}

func (s *kmsCryptoKeyId) terraformId() string {
	return fmt.Sprintf("%s/%s", s.KeyRingId.terraformId(), s.Name)
}

func validateKmsCryptoKeyRotationPeriod(value interface{}, _ string) (ws []string, errors []error) {
	period := value.(string)
	pattern := regexp.MustCompile(`^([0-9.]*\d)s$`)
	match := pattern.FindStringSubmatch(period)

	if len(match) == 0 {
		errors = append(errors, fmt.Errorf("Invalid rotation period format: %s", period))
		// Cannot continue to validate because we cannot extract a number.
		return
	}

	number := match[1]
	seconds, err := strconv.ParseFloat(number, 64)

	if err != nil {
		errors = append(errors, err)
	} else {
		if seconds < 86400.0 {
			errors = append(errors, fmt.Errorf("Rotation period must be greater than one day"))
		}

		parts := strings.Split(number, ".")

		if len(parts) > 1 && len(parts[1]) > 9 {
			errors = append(errors, fmt.Errorf("Rotation period cannot have more than 9 fractional digits"))
		}
	}

	return
}

func kmsCryptoKeyNextRotation(now time.Time, period string) (result string, err error) {
	var duration time.Duration

	duration, err = time.ParseDuration(period)

	if err == nil {
		result = now.UTC().Add(duration).Format(time.RFC3339Nano)
	}

	return
}

func parseKmsCryptoKeyId(id string, config *Config) (*kmsCryptoKeyId, error) {
	parts := strings.Split(id, "/")

	cryptoKeyIdRegex := regexp.MustCompile("^(" + ProjectRegex + ")/([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})/([a-zA-Z0-9_-]{1,63})$")
	cryptoKeyIdWithoutProjectRegex := regexp.MustCompile("^([a-z0-9-])+/([a-zA-Z0-9_-]{1,63})/([a-zA-Z0-9_-]{1,63})$")
	cryptoKeyRelativeLinkRegex := regexp.MustCompile("^projects/(" + ProjectRegex + ")/locations/([a-z0-9-]+)/keyRings/([a-zA-Z0-9_-]{1,63})/cryptoKeys/([a-zA-Z0-9_-]{1,63})$")

	if cryptoKeyIdRegex.MatchString(id) {
		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  parts[0],
				Location: parts[1],
				Name:     parts[2],
			},
			Name: parts[3],
		}, nil
	}

	if cryptoKeyIdWithoutProjectRegex.MatchString(id) {
		if config.Project == "" {
			return nil, fmt.Errorf("The default project for the provider must be set when using the `{location}/{keyRingName}/{cryptoKeyName}` id format.")
		}

		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  config.Project,
				Location: parts[0],
				Name:     parts[1],
			},
			Name: parts[2],
		}, nil
	}

	if parts := cryptoKeyRelativeLinkRegex.FindStringSubmatch(id); parts != nil {
		return &kmsCryptoKeyId{
			KeyRingId: kmsKeyRingId{
				Project:  parts[1],
				Location: parts[2],
				Name:     parts[3],
			},
			Name: parts[4],
		}, nil
	}
	return nil, fmt.Errorf("Invalid CryptoKey id format, expecting `{projectId}/{locationId}/{KeyringName}/{cryptoKeyName}` or `{locationId}/{keyRingName}/{cryptoKeyName}, got id: %s`", id)
}

func clearCryptoKeyVersions(cryptoKeyId *kmsCryptoKeyId, userAgent string, config *Config) error {
	versionsClient := config.NewKmsClient(userAgent).Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions

	listCall := versionsClient.List(cryptoKeyId.cryptoKeyId())
	if config.UserProjectOverride {
		listCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
	}
	versionsResponse, err := listCall.Do()

	if err != nil {
		return err
	}

	for _, version := range versionsResponse.CryptoKeyVersions {
		// skip the versions that have been destroyed earlier
		if version.State == "ENABLED" {
			request := &cloudkms.DestroyCryptoKeyVersionRequest{}
			destroyCall := versionsClient.Destroy(version.Name, request)
			if config.UserProjectOverride {
				destroyCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
			}
			_, err = destroyCall.Do()

			if err != nil {
				return err
			}
		}
	}

	return nil
}

func disableCryptoKeyRotation(cryptoKeyId *kmsCryptoKeyId, userAgent string, config *Config) error {
	keyClient := config.NewKmsClient(userAgent).Projects.Locations.KeyRings.CryptoKeys
	patchCall := keyClient.Patch(cryptoKeyId.cryptoKeyId(), &cloudkms.CryptoKey{
		NullFields: []string{"rotationPeriod", "nextRotationTime"},
	}).
		UpdateMask("rotationPeriod,nextRotationTime")
	if config.UserProjectOverride {
		patchCall.Header().Set("X-Goog-User-Project", cryptoKeyId.KeyRingId.Project)
	}
	_, err := patchCall.Do()

	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed 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 framework

import (
	"github.com/kubernetes-sigs/kube-batch/pkg/scheduler/api"
)

func (ssn *Session) AddJobOrderFn(name string, cf api.CompareFn) {
	ssn.jobOrderFns[name] = cf
}

func (ssn *Session) AddQueueOrderFn(name string, qf api.CompareFn) {
	ssn.queueOrderFns[name] = qf
}

func (ssn *Session) AddTaskOrderFn(name string, cf api.CompareFn) {
	ssn.taskOrderFns[name] = cf
}

func (ssn *Session) AddPreemptableFn(name string, cf api.EvictableFn) {
	ssn.preemptableFns[name] = cf
}

func (ssn *Session) AddReclaimableFn(name string, rf api.EvictableFn) {
	ssn.reclaimableFns[name] = rf
}

func (ssn *Session) AddJobReadyFn(name string, vf api.ValidateFn) {
	ssn.jobReadyFns[name] = vf
}

func (ssn *Session) AddPredicateFn(name string, pf api.PredicateFn) {
	ssn.predicateFns[name] = pf
}

func (ssn *Session) AddNodeOrderFn(name string, pf api.NodeOrderFn) {
	ssn.nodeOrderFns[name] = pf
}

func (ssn *Session) AddOverusedFn(name string, fn api.ValidateFn) {
	ssn.overusedFns[name] = fn
}

func (ssn *Session) AddJobValidFn(name string, fn api.ValidateExFn) {
	ssn.jobValidFns[name] = fn
}

func (ssn *Session) Reclaimable(reclaimer *api.TaskInfo, reclaimees []*api.TaskInfo) []*api.TaskInfo {
	var victims []*api.TaskInfo
	var init bool

	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.ReclaimableDisabled {
				continue
			}
			rf, found := ssn.reclaimableFns[plugin.Name]
			if !found {
				continue
			}
			candidates := rf(reclaimer, reclaimees)
			if !init {
				victims = candidates
				init = true
			} else {
				var intersection []*api.TaskInfo
				// Get intersection of victims and candidates.
				for _, v := range victims {
					for _, c := range candidates {
						if v.UID == c.UID {
							intersection = append(intersection, v)
						}
					}
				}

				// Update victims to intersection
				victims = intersection
			}
		}
		// Plugins in this tier made decision if victims is not nil
		if victims != nil {
			return victims
		}
	}

	return victims
}

func (ssn *Session) Preemptable(preemptor *api.TaskInfo, preemptees []*api.TaskInfo) []*api.TaskInfo {
	var victims []*api.TaskInfo
	var init bool

	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.PreemptableDisabled {
				continue
			}

			pf, found := ssn.preemptableFns[plugin.Name]
			if !found {
				continue
			}
			candidates := pf(preemptor, preemptees)
			if !init {
				victims = candidates
				init = true
			} else {
				var intersection []*api.TaskInfo
				// Get intersection of victims and candidates.
				for _, v := range victims {
					for _, c := range candidates {
						if v.UID == c.UID {
							intersection = append(intersection, v)
						}
					}
				}

				// Update victims to intersection
				victims = intersection
			}
		}
		// Plugins in this tier made decision if victims is not nil
		if victims != nil {
			return victims
		}
	}

	return victims
}

func (ssn *Session) Overused(queue *api.QueueInfo) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			of, found := ssn.overusedFns[plugin.Name]
			if !found {
				continue
			}
			if of(queue) {
				return true
			}
		}
	}

	return false
}

func (ssn *Session) JobReady(obj interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.JobReadyDisabled {
				continue
			}
			jrf, found := ssn.jobReadyFns[plugin.Name]
			if !found {
				continue
			}

			if !jrf(obj) {
				return false
			}
		}
	}

	return true
}

func (ssn *Session) JobValid(obj interface{}) *api.ValidateResult {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			jrf, found := ssn.jobValidFns[plugin.Name]
			if !found {
				continue
			}

			if vr := jrf(obj); vr != nil && !vr.Pass {
				return vr
			}

		}
	}

	return nil
}

func (ssn *Session) JobOrderFn(l, r interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.JobOrderDisabled {
				continue
			}
			jof, found := ssn.jobOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := jof(l, r); j != 0 {
				return j < 0
			}
		}
	}

	// If no job order funcs, order job by UID.
	lv := l.(*api.JobInfo)
	rv := r.(*api.JobInfo)

	if lv.CreationTimestamp.Equal(&rv.CreationTimestamp) {
		return lv.UID < rv.UID
	}

	return lv.CreationTimestamp.Before(&rv.CreationTimestamp)
}

func (ssn *Session) QueueOrderFn(l, r interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.QueueOrderDisabled {
				continue
			}
			qof, found := ssn.queueOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := qof(l, r); j != 0 {
				return j < 0
			}

		}
	}

	// If no queue order funcs, order queue by UID.
	lv := l.(*api.QueueInfo)
	rv := r.(*api.QueueInfo)

	return lv.UID < rv.UID
}

func (ssn *Session) TaskCompareFns(l, r interface{}) int {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.TaskOrderDisabled {
				continue
			}
			tof, found := ssn.taskOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := tof(l, r); j != 0 {
				return j
			}
		}
	}

	return 0
}

func (ssn *Session) TaskOrderFn(l, r interface{}) bool {
	if res := ssn.TaskCompareFns(l, r); res != 0 {
		return res < 0
	}

	// If no task order funcs, order task by CreationTimestamp first, then by UID.
	lv := l.(*api.TaskInfo)
	rv := r.(*api.TaskInfo)
	if lv.Pod.CreationTimestamp.Equal(&rv.Pod.CreationTimestamp) {
		return lv.UID < rv.UID
	} else {
		return lv.Pod.CreationTimestamp.Before(&rv.Pod.CreationTimestamp)
	}
}

func (ssn *Session) PredicateFn(task *api.TaskInfo, node *api.NodeInfo) error {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.PredicateDisabled {
				continue
			}
			pfn, found := ssn.predicateFns[plugin.Name]
			if !found {
				continue
			}
			err := pfn(task, node)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (ssn *Session) NodeOrderFn(task *api.TaskInfo, node *api.NodeInfo) (int, error) {
	priorityScore := 0
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.NodeOrderDisabled {
				continue
			}
			pfn, found := ssn.nodeOrderFns[plugin.Name]
			if !found {
				continue
			}
			score, err := pfn(task, node)
			if err != nil {
				return 0, err
			} else {
				priorityScore = priorityScore + score
			}
		}
	}
	return priorityScore, nil
}
<commit_msg>Order queue by CreationTimestamp first, then by UID<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed 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 framework

import (
	"github.com/kubernetes-sigs/kube-batch/pkg/scheduler/api"
)

func (ssn *Session) AddJobOrderFn(name string, cf api.CompareFn) {
	ssn.jobOrderFns[name] = cf
}

func (ssn *Session) AddQueueOrderFn(name string, qf api.CompareFn) {
	ssn.queueOrderFns[name] = qf
}

func (ssn *Session) AddTaskOrderFn(name string, cf api.CompareFn) {
	ssn.taskOrderFns[name] = cf
}

func (ssn *Session) AddPreemptableFn(name string, cf api.EvictableFn) {
	ssn.preemptableFns[name] = cf
}

func (ssn *Session) AddReclaimableFn(name string, rf api.EvictableFn) {
	ssn.reclaimableFns[name] = rf
}

func (ssn *Session) AddJobReadyFn(name string, vf api.ValidateFn) {
	ssn.jobReadyFns[name] = vf
}

func (ssn *Session) AddPredicateFn(name string, pf api.PredicateFn) {
	ssn.predicateFns[name] = pf
}

func (ssn *Session) AddNodeOrderFn(name string, pf api.NodeOrderFn) {
	ssn.nodeOrderFns[name] = pf
}

func (ssn *Session) AddOverusedFn(name string, fn api.ValidateFn) {
	ssn.overusedFns[name] = fn
}

func (ssn *Session) AddJobValidFn(name string, fn api.ValidateExFn) {
	ssn.jobValidFns[name] = fn
}

func (ssn *Session) Reclaimable(reclaimer *api.TaskInfo, reclaimees []*api.TaskInfo) []*api.TaskInfo {
	var victims []*api.TaskInfo
	var init bool

	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.ReclaimableDisabled {
				continue
			}
			rf, found := ssn.reclaimableFns[plugin.Name]
			if !found {
				continue
			}
			candidates := rf(reclaimer, reclaimees)
			if !init {
				victims = candidates
				init = true
			} else {
				var intersection []*api.TaskInfo
				// Get intersection of victims and candidates.
				for _, v := range victims {
					for _, c := range candidates {
						if v.UID == c.UID {
							intersection = append(intersection, v)
						}
					}
				}

				// Update victims to intersection
				victims = intersection
			}
		}
		// Plugins in this tier made decision if victims is not nil
		if victims != nil {
			return victims
		}
	}

	return victims
}

func (ssn *Session) Preemptable(preemptor *api.TaskInfo, preemptees []*api.TaskInfo) []*api.TaskInfo {
	var victims []*api.TaskInfo
	var init bool

	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.PreemptableDisabled {
				continue
			}

			pf, found := ssn.preemptableFns[plugin.Name]
			if !found {
				continue
			}
			candidates := pf(preemptor, preemptees)
			if !init {
				victims = candidates
				init = true
			} else {
				var intersection []*api.TaskInfo
				// Get intersection of victims and candidates.
				for _, v := range victims {
					for _, c := range candidates {
						if v.UID == c.UID {
							intersection = append(intersection, v)
						}
					}
				}

				// Update victims to intersection
				victims = intersection
			}
		}
		// Plugins in this tier made decision if victims is not nil
		if victims != nil {
			return victims
		}
	}

	return victims
}

func (ssn *Session) Overused(queue *api.QueueInfo) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			of, found := ssn.overusedFns[plugin.Name]
			if !found {
				continue
			}
			if of(queue) {
				return true
			}
		}
	}

	return false
}

func (ssn *Session) JobReady(obj interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.JobReadyDisabled {
				continue
			}
			jrf, found := ssn.jobReadyFns[plugin.Name]
			if !found {
				continue
			}

			if !jrf(obj) {
				return false
			}
		}
	}

	return true
}

func (ssn *Session) JobValid(obj interface{}) *api.ValidateResult {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			jrf, found := ssn.jobValidFns[plugin.Name]
			if !found {
				continue
			}

			if vr := jrf(obj); vr != nil && !vr.Pass {
				return vr
			}

		}
	}

	return nil
}

func (ssn *Session) JobOrderFn(l, r interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.JobOrderDisabled {
				continue
			}
			jof, found := ssn.jobOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := jof(l, r); j != 0 {
				return j < 0
			}
		}
	}

	// If no job order funcs, order job by UID.
	lv := l.(*api.JobInfo)
	rv := r.(*api.JobInfo)

	if lv.CreationTimestamp.Equal(&rv.CreationTimestamp) {
		return lv.UID < rv.UID
	}

	return lv.CreationTimestamp.Before(&rv.CreationTimestamp)
}

func (ssn *Session) QueueOrderFn(l, r interface{}) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.QueueOrderDisabled {
				continue
			}
			qof, found := ssn.queueOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := qof(l, r); j != 0 {
				return j < 0
			}

		}
	}

	// If no queue order funcs, order queue by CreationTimestamp first, then by UID.
	lv := l.(*api.QueueInfo)
	rv := r.(*api.QueueInfo)
	if lv.Queue.CreationTimestamp.Equal(&rv.Queue.CreationTimestamp) {
		return lv.UID < rv.UID
	} else {
		return lv.Queue.CreationTimestamp.Before(&rv.Queue.CreationTimestamp)
	}
}

func (ssn *Session) TaskCompareFns(l, r interface{}) int {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.TaskOrderDisabled {
				continue
			}
			tof, found := ssn.taskOrderFns[plugin.Name]
			if !found {
				continue
			}
			if j := tof(l, r); j != 0 {
				return j
			}
		}
	}

	return 0
}

func (ssn *Session) TaskOrderFn(l, r interface{}) bool {
	if res := ssn.TaskCompareFns(l, r); res != 0 {
		return res < 0
	}

	// If no task order funcs, order task by CreationTimestamp first, then by UID.
	lv := l.(*api.TaskInfo)
	rv := r.(*api.TaskInfo)
	if lv.Pod.CreationTimestamp.Equal(&rv.Pod.CreationTimestamp) {
		return lv.UID < rv.UID
	} else {
		return lv.Pod.CreationTimestamp.Before(&rv.Pod.CreationTimestamp)
	}
}

func (ssn *Session) PredicateFn(task *api.TaskInfo, node *api.NodeInfo) error {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.PredicateDisabled {
				continue
			}
			pfn, found := ssn.predicateFns[plugin.Name]
			if !found {
				continue
			}
			err := pfn(task, node)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (ssn *Session) NodeOrderFn(task *api.TaskInfo, node *api.NodeInfo) (int, error) {
	priorityScore := 0
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.NodeOrderDisabled {
				continue
			}
			pfn, found := ssn.nodeOrderFns[plugin.Name]
			if !found {
				continue
			}
			score, err := pfn(task, node)
			if err != nil {
				return 0, err
			} else {
				priorityScore = priorityScore + score
			}
		}
	}
	return priorityScore, nil
}
<|endoftext|>


<commit_before>package notifiers

import (
	"bytes"
	"fmt"
	"io"
	"mime/multipart"
	"os"
	"strconv"
	"strings"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/alerting"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	alerting.RegisterNotifier(&alerting.NotifierPlugin{
		Type:        "discord",
		Name:        "Discord",
		Description: "Sends notifications to Discord",
		Factory:     newDiscordNotifier,
		Heading:     "Discord settings",
		Options: []alerting.NotifierOption{
			{
				Label:        "Message Content",
				Description:  "Mention a group using @ or a user using <@ID> when notifying in a channel",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				PropertyName: "content",
			},
			{
				Label:        "Webhook URL",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Placeholder:  "Discord webhook URL",
				PropertyName: "url",
				Required:     true,
			},
		},
	})
}

func newDiscordNotifier(model *models.AlertNotification) (alerting.Notifier, error) {
	content := model.Settings.Get("content").MustString()
	url := model.Settings.Get("url").MustString()
	if url == "" {
		return nil, alerting.ValidationError{Reason: "Could not find webhook url property in settings"}
	}

	return &DiscordNotifier{
		NotifierBase: NewNotifierBase(model),
		Content:      content,
		WebhookURL:   url,
		log:          log.New("alerting.notifier.discord"),
	}, nil
}

// DiscordNotifier is responsible for sending alert
// notifications to discord.
type DiscordNotifier struct {
	NotifierBase
	Content    string
	WebhookURL string
	log        log.Logger
}

// Notify send an alert notification to Discord.
func (dn *DiscordNotifier) Notify(evalContext *alerting.EvalContext) error {
	dn.log.Info("Sending alert notification to", "webhook_url", dn.WebhookURL)

	ruleURL, err := evalContext.GetRuleURL()
	if err != nil {
		dn.log.Error("Failed get rule link", "error", err)
		return err
	}

	bodyJSON := simplejson.New()
	bodyJSON.Set("username", "Grafana")

	if dn.Content != "" {
		bodyJSON.Set("content", dn.Content)
	}

	fields := make([]map[string]interface{}, 0)

	for _, evt := range evalContext.EvalMatches {
		fields = append(fields, map[string]interface{}{
			"name":   evt.Metric,
			"value":  evt.Value.FullString(),
			"inline": true,
		})
	}

	footer := map[string]interface{}{
		"text":     "Grafana v" + setting.BuildVersion,
		"icon_url": "https://grafana.com/assets/img/fav32.png",
	}

	color, _ := strconv.ParseInt(strings.TrimLeft(evalContext.GetStateModel().Color, "#"), 16, 0)

	embed := simplejson.New()
	embed.Set("title", evalContext.GetNotificationTitle())
	// Discord takes integer for color
	embed.Set("color", color)
	embed.Set("url", ruleURL)
	embed.Set("description", evalContext.Rule.Message)
	embed.Set("type", "rich")
	embed.Set("fields", fields)
	embed.Set("footer", footer)

	var image map[string]interface{}
	var embeddedImage = false

	if dn.NeedsImage() {
		if evalContext.ImagePublicURL != "" {
			image = map[string]interface{}{
				"url": evalContext.ImagePublicURL,
			}
			embed.Set("image", image)
		} else {
			image = map[string]interface{}{
				"url": "attachment://graph.png",
			}
			embed.Set("image", image)
			embeddedImage = true
		}
	}

	bodyJSON.Set("embeds", []interface{}{embed})

	json, _ := bodyJSON.MarshalJSON()

	cmd := &models.SendWebhookSync{
		Url:         dn.WebhookURL,
		HttpMethod:  "POST",
		ContentType: "application/json",
	}

	if !embeddedImage {
		cmd.Body = string(json)
	} else {
		err := dn.embedImage(cmd, evalContext.ImageOnDiskPath, json)
		if err != nil {
			dn.log.Error("failed to embed image", "error", err)
			return err
		}
	}

	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		dn.log.Error("Failed to send notification to Discord", "error", err)
		return err
	}

	return nil
}

func (dn *DiscordNotifier) embedImage(cmd *models.SendWebhookSync, imagePath string, existingJSONBody []byte) error {
	// nolint:gosec
	// We can ignore the gosec G304 warning on this one because `imagePath` comes
	// from the alert `evalContext` that generates the images.
	f, err := os.Open(imagePath)
	if err != nil {
		if os.IsNotExist(err) {
			cmd.Body = string(existingJSONBody)
			return nil
		}
		if !os.IsNotExist(err) {
			return err
		}
	}
	defer func() {
		if err := f.Close(); err != nil {
			dn.log.Warn("Failed to close file", "path", imagePath, "err", err)
		}
	}()

	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	defer func() {
		if err := w.Close(); err != nil {
			// Should be OK since we already close it on non-error path
			dn.log.Warn("Failed to close multipart writer", "err", err)
		}
	}()
	fw, err := w.CreateFormField("payload_json")
	if err != nil {
		return err
	}

	if _, err = fw.Write([]byte(string(existingJSONBody))); err != nil {
		return err
	}

	fw, err = w.CreateFormFile("file", "graph.png")
	if err != nil {
		return err
	}

	if _, err = io.Copy(fw, f); err != nil {
		return err
	}

	if err := w.Close(); err != nil {
		return fmt.Errorf("failed to close multipart writer: %w", err)
	}

	cmd.Body = b.String()
	cmd.ContentType = w.FormDataContentType()

	return nil
}
<commit_msg>Alerting: Fix bug in Discord for when name for metric value is absent (#31257)<commit_after>package notifiers

import (
	"bytes"
	"fmt"
	"io"
	"mime/multipart"
	"os"
	"strconv"
	"strings"

	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/services/alerting"
	"github.com/grafana/grafana/pkg/setting"
)

func init() {
	alerting.RegisterNotifier(&alerting.NotifierPlugin{
		Type:        "discord",
		Name:        "Discord",
		Description: "Sends notifications to Discord",
		Factory:     newDiscordNotifier,
		Heading:     "Discord settings",
		Options: []alerting.NotifierOption{
			{
				Label:        "Message Content",
				Description:  "Mention a group using @ or a user using <@ID> when notifying in a channel",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				PropertyName: "content",
			},
			{
				Label:        "Webhook URL",
				Element:      alerting.ElementTypeInput,
				InputType:    alerting.InputTypeText,
				Placeholder:  "Discord webhook URL",
				PropertyName: "url",
				Required:     true,
			},
		},
	})
}

func newDiscordNotifier(model *models.AlertNotification) (alerting.Notifier, error) {
	content := model.Settings.Get("content").MustString()
	url := model.Settings.Get("url").MustString()
	if url == "" {
		return nil, alerting.ValidationError{Reason: "Could not find webhook url property in settings"}
	}

	return &DiscordNotifier{
		NotifierBase: NewNotifierBase(model),
		Content:      content,
		WebhookURL:   url,
		log:          log.New("alerting.notifier.discord"),
	}, nil
}

// DiscordNotifier is responsible for sending alert
// notifications to discord.
type DiscordNotifier struct {
	NotifierBase
	Content    string
	WebhookURL string
	log        log.Logger
}

// Notify send an alert notification to Discord.
func (dn *DiscordNotifier) Notify(evalContext *alerting.EvalContext) error {
	dn.log.Info("Sending alert notification to", "webhook_url", dn.WebhookURL)

	ruleURL, err := evalContext.GetRuleURL()
	if err != nil {
		dn.log.Error("Failed get rule link", "error", err)
		return err
	}

	bodyJSON := simplejson.New()
	bodyJSON.Set("username", "Grafana")

	if dn.Content != "" {
		bodyJSON.Set("content", dn.Content)
	}

	fields := make([]map[string]interface{}, 0)

	for _, evt := range evalContext.EvalMatches {
		fields = append(fields, map[string]interface{}{
			// Discord uniquely does not send the alert if the metric field is empty,
			// which it can be in some cases
			"name":   notEmpty(evt.Metric),
			"value":  evt.Value.FullString(),
			"inline": true,
		})
	}

	footer := map[string]interface{}{
		"text":     "Grafana v" + setting.BuildVersion,
		"icon_url": "https://grafana.com/assets/img/fav32.png",
	}

	color, _ := strconv.ParseInt(strings.TrimLeft(evalContext.GetStateModel().Color, "#"), 16, 0)

	embed := simplejson.New()
	embed.Set("title", evalContext.GetNotificationTitle())
	// Discord takes integer for color
	embed.Set("color", color)
	embed.Set("url", ruleURL)
	embed.Set("description", evalContext.Rule.Message)
	embed.Set("type", "rich")
	embed.Set("fields", fields)
	embed.Set("footer", footer)

	var image map[string]interface{}
	var embeddedImage = false

	if dn.NeedsImage() {
		if evalContext.ImagePublicURL != "" {
			image = map[string]interface{}{
				"url": evalContext.ImagePublicURL,
			}
			embed.Set("image", image)
		} else {
			image = map[string]interface{}{
				"url": "attachment://graph.png",
			}
			embed.Set("image", image)
			embeddedImage = true
		}
	}

	bodyJSON.Set("embeds", []interface{}{embed})

	json, _ := bodyJSON.MarshalJSON()

	cmd := &models.SendWebhookSync{
		Url:         dn.WebhookURL,
		HttpMethod:  "POST",
		ContentType: "application/json",
	}

	if !embeddedImage {
		cmd.Body = string(json)
	} else {
		err := dn.embedImage(cmd, evalContext.ImageOnDiskPath, json)
		if err != nil {
			dn.log.Error("failed to embed image", "error", err)
			return err
		}
	}

	if err := bus.DispatchCtx(evalContext.Ctx, cmd); err != nil {
		dn.log.Error("Failed to send notification to Discord", "error", err)
		return err
	}

	return nil
}

func (dn *DiscordNotifier) embedImage(cmd *models.SendWebhookSync, imagePath string, existingJSONBody []byte) error {
	// nolint:gosec
	// We can ignore the gosec G304 warning on this one because `imagePath` comes
	// from the alert `evalContext` that generates the images.
	f, err := os.Open(imagePath)
	if err != nil {
		if os.IsNotExist(err) {
			cmd.Body = string(existingJSONBody)
			return nil
		}
		if !os.IsNotExist(err) {
			return err
		}
	}
	defer func() {
		if err := f.Close(); err != nil {
			dn.log.Warn("Failed to close file", "path", imagePath, "err", err)
		}
	}()

	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	defer func() {
		if err := w.Close(); err != nil {
			// Should be OK since we already close it on non-error path
			dn.log.Warn("Failed to close multipart writer", "err", err)
		}
	}()
	fw, err := w.CreateFormField("payload_json")
	if err != nil {
		return err
	}

	if _, err = fw.Write([]byte(string(existingJSONBody))); err != nil {
		return err
	}

	fw, err = w.CreateFormFile("file", "graph.png")
	if err != nil {
		return err
	}

	if _, err = io.Copy(fw, f); err != nil {
		return err
	}

	if err := w.Close(); err != nil {
		return fmt.Errorf("failed to close multipart writer: %w", err)
	}

	cmd.Body = b.String()
	cmd.ContentType = w.FormDataContentType()

	return nil
}

func notEmpty(metric string) string {
	if metric == "" {
		return "<NO_METRIC_NAME>"
	}

	return metric
}
<|endoftext|>


<commit_before>package shipper

import (
	"context"
	"fmt"
	"strings"

	"github.com/go-kit/kit/log/level"

	"github.com/cortexproject/cortex/pkg/chunk"
	cortex_util "github.com/cortexproject/cortex/pkg/util"

	"github.com/grafana/loki/pkg/storage/stores/util"
)

const (
	delimiter = "/"
)

type boltDBShipperTableClient struct {
	objectClient chunk.ObjectClient
}

func NewBoltDBShipperTableClient(objectClient chunk.ObjectClient) chunk.TableClient {
	return &boltDBShipperTableClient{util.NewPrefixedObjectClient(objectClient, StorageKeyPrefix)}
}

func (b *boltDBShipperTableClient) ListTables(ctx context.Context) ([]string, error) {
	_, dirs, err := b.objectClient.List(ctx, "", delimiter)
	if err != nil {
		return nil, err
	}

	tables := make([]string, len(dirs))
	for i, dir := range dirs {
		tables[i] = strings.TrimSuffix(string(dir), delimiter)
	}

	return tables, nil
}

func (b *boltDBShipperTableClient) CreateTable(ctx context.Context, desc chunk.TableDesc) error {
	return nil
}

func (b *boltDBShipperTableClient) Stop() {
	b.objectClient.Stop()
}

func (b *boltDBShipperTableClient) DeleteTable(ctx context.Context, name string) error {
	objects, dirs, err := b.objectClient.List(ctx, name, delimiter)
	if err != nil {
		return err
	}

	if len(dirs) != 0 {
		level.Error(cortex_util.Logger).Log("msg", fmt.Sprintf("unexpected directories in %s folder, not touching them", name), "directories", fmt.Sprint(dirs))
	}

	for _, object := range objects {
		err := b.objectClient.DeleteObject(ctx, object.Key)
		if err != nil {
			return err
		}
	}

	return nil
}

func (b *boltDBShipperTableClient) DescribeTable(ctx context.Context, name string) (desc chunk.TableDesc, isActive bool, err error) {
	return chunk.TableDesc{
		Name: name,
	}, true, nil
}

func (b *boltDBShipperTableClient) UpdateTable(ctx context.Context, current, expected chunk.TableDesc) error {
	return nil
}
<commit_msg>fix table deletion in table client for boltdb-shipper (#2960)<commit_after>package shipper

import (
	"context"
	"fmt"
	"strings"

	"github.com/go-kit/kit/log/level"

	"github.com/cortexproject/cortex/pkg/chunk"
	cortex_util "github.com/cortexproject/cortex/pkg/util"

	"github.com/grafana/loki/pkg/storage/stores/util"
)

const (
	delimiter = "/"
)

type boltDBShipperTableClient struct {
	objectClient chunk.ObjectClient
}

func NewBoltDBShipperTableClient(objectClient chunk.ObjectClient) chunk.TableClient {
	return &boltDBShipperTableClient{util.NewPrefixedObjectClient(objectClient, StorageKeyPrefix)}
}

func (b *boltDBShipperTableClient) ListTables(ctx context.Context) ([]string, error) {
	_, dirs, err := b.objectClient.List(ctx, "", delimiter)
	if err != nil {
		return nil, err
	}

	tables := make([]string, len(dirs))
	for i, dir := range dirs {
		tables[i] = strings.TrimSuffix(string(dir), delimiter)
	}

	return tables, nil
}

func (b *boltDBShipperTableClient) CreateTable(ctx context.Context, desc chunk.TableDesc) error {
	return nil
}

func (b *boltDBShipperTableClient) Stop() {
	b.objectClient.Stop()
}

func (b *boltDBShipperTableClient) DeleteTable(ctx context.Context, name string) error {
	objects, dirs, err := b.objectClient.List(ctx, name+delimiter, delimiter)
	if err != nil {
		return err
	}

	if len(dirs) != 0 {
		level.Error(cortex_util.Logger).Log("msg", fmt.Sprintf("unexpected directories in %s folder, not touching them", name), "directories", fmt.Sprint(dirs))
	}

	for _, object := range objects {
		err := b.objectClient.DeleteObject(ctx, object.Key)
		if err != nil {
			return err
		}
	}

	return nil
}

func (b *boltDBShipperTableClient) DescribeTable(ctx context.Context, name string) (desc chunk.TableDesc, isActive bool, err error) {
	return chunk.TableDesc{
		Name: name,
	}, true, nil
}

func (b *boltDBShipperTableClient) UpdateTable(ctx context.Context, current, expected chunk.TableDesc) error {
	return nil
}
<|endoftext|>