blob: 310f4ec2b50a579679894b7d0a5d31f4c2245f3c [file] [log] [blame]
// Copyright 2020 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 invocations
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"sort"
"cloud.google.com/go/spanner"
"go.chromium.org/luci/resultdb/internal/spanutil"
"go.chromium.org/luci/resultdb/pbutil"
)
// ID can convert an invocation id to various formats.
type ID string
// ToSpanner implements span.Value.
func (id ID) ToSpanner() interface{} {
return id.RowID()
}
// SpannerPtr implements span.Ptr.
func (id *ID) SpannerPtr(b *spanutil.Buffer) interface{} {
return &b.NullString
}
// FromSpanner implements span.Ptr.
func (id *ID) FromSpanner(b *spanutil.Buffer) error {
*id = ""
if b.NullString.Valid {
*id = IDFromRowID(b.NullString.StringVal)
}
return nil
}
// MustParseName converts an invocation name to an ID.
// Panics if the name is invalid. Useful for situations when name was already
// validated.
func MustParseName(name string) ID {
id, err := pbutil.ParseInvocationName(name)
if err != nil {
panic(err)
}
return ID(id)
}
// IDFromRowID converts a Spanner-level row ID to an ID.
func IDFromRowID(rowID string) ID {
return ID(stripHashPrefix(rowID))
}
// Name returns an invocation name.
func (id ID) Name() string {
return pbutil.InvocationName(string(id))
}
// RowID returns an invocation ID used in spanner rows.
// If id is empty, returns "".
func (id ID) RowID() string {
if id == "" {
return ""
}
return prefixWithHash(string(id))
}
// Key returns a invocation spanner key.
func (id ID) Key(suffix ...interface{}) spanner.Key {
ret := make(spanner.Key, 1+len(suffix))
ret[0] = id.RowID()
copy(ret[1:], suffix)
return ret
}
// IDSet is an unordered set of invocation ids.
type IDSet map[ID]struct{}
// NewIDSet creates an IDSet from members.
func NewIDSet(ids ...ID) IDSet {
ret := make(IDSet, len(ids))
for _, id := range ids {
ret.Add(id)
}
return ret
}
// Add adds id to the set.
func (s IDSet) Add(id ID) {
s[id] = struct{}{}
}
// Union adds other ids.
func (s IDSet) Union(other IDSet) {
for id := range other {
s.Add(id)
}
}
// RemoveAll removes any ids present in other.
func (s IDSet) RemoveAll(other IDSet) {
if len(s) > 0 {
for id := range other {
s.Remove(id)
}
}
}
// Remove removes id from the set if it was present.
func (s IDSet) Remove(id ID) {
delete(s, id)
}
// Has returns true if id is in the set.
func (s IDSet) Has(id ID) bool {
_, ok := s[id]
return ok
}
// String implements fmt.Stringer.
func (s IDSet) String() string {
strs := make([]string, 0, len(s))
for id := range s {
strs = append(strs, string(id))
}
sort.Strings(strs)
return fmt.Sprintf("%q", strs)
}
// Keys returns a spanner.KeySet.
func (s IDSet) Keys(suffix ...interface{}) spanner.KeySet {
ret := spanner.KeySets()
for id := range s {
ret = spanner.KeySets(id.Key(suffix...), ret)
}
return ret
}
// ToSpanner implements span.Value.
func (s IDSet) ToSpanner() interface{} {
ret := make([]string, 0, len(s))
for id := range s {
ret = append(ret, id.RowID())
}
sort.Strings(ret)
return ret
}
// SpannerPtr implements span.Ptr.
func (s *IDSet) SpannerPtr(b *spanutil.Buffer) interface{} {
return &b.StringSlice
}
// FromSpanner implements span.Ptr.
func (s *IDSet) FromSpanner(b *spanutil.Buffer) error {
*s = make(IDSet, len(b.StringSlice))
for _, rowID := range b.StringSlice {
s.Add(IDFromRowID(rowID))
}
return nil
}
// ParseNames converts invocation names to IDSet.
func ParseNames(names []string) (IDSet, error) {
ids := make(IDSet, len(names))
for _, name := range names {
id, err := pbutil.ParseInvocationName(name)
if err != nil {
return nil, err
}
ids.Add(ID(id))
}
return ids, nil
}
// MustParseNames converts invocation names to IDSet.
// Panics if a name is invalid. Useful for situations when names were already
// validated.
func MustParseNames(names []string) IDSet {
ids, err := ParseNames(names)
if err != nil {
panic(err)
}
return ids
}
// Names returns a sorted slice of invocation names.
func (s IDSet) Names() []string {
names := make([]string, 0, len(s))
for id := range s {
names = append(names, id.Name())
}
sort.Strings(names)
return names
}
// Batches splits s into batches.
// The batches are sorted by RowID(), such that interval (minRowID, maxRowID)
// of each batch does not overlap with any other batch.
//
// The size of batch is hardcoded 50, because that's the maximum parallelism
// we get from Cloud Spanner.
func (s IDSet) Batches() []IDSet {
return s.batches(50)
}
func (s IDSet) batches(size int) []IDSet {
ids := s.SortByRowID()
batches := make([]IDSet, 0, 1+len(ids)/size)
for len(ids) > 0 {
batchSize := size
if batchSize > len(ids) {
batchSize = len(ids)
}
batches = append(batches, NewIDSet(ids[:batchSize]...))
ids = ids[batchSize:]
}
return batches
}
// SortByRowID returns IDs in the set sorted by row id.
func (s IDSet) SortByRowID() []ID {
rowIDs := make([]string, 0, len(s))
for id := range s {
rowIDs = append(rowIDs, id.RowID())
}
sort.Strings(rowIDs)
ret := make([]ID, len(rowIDs))
for i, rowID := range rowIDs {
ret[i] = ID(stripHashPrefix(rowID))
}
return ret
}
// hashPrefixBytes is the number of bytes of sha256 to prepend to a PK
// to achieve even distribution.
const hashPrefixBytes = 4
func prefixWithHash(s string) string {
h := sha256.Sum256([]byte(s))
prefix := hex.EncodeToString(h[:hashPrefixBytes])
return fmt.Sprintf("%s:%s", prefix, s)
}
func stripHashPrefix(s string) string {
expectedPrefixLen := hex.EncodedLen(hashPrefixBytes) + 1 // +1 for separator
if len(s) < expectedPrefixLen {
panic(fmt.Sprintf("%q is too short", s))
}
return s[expectedPrefixLen:]
}