spanset.go

// Copyright 2017 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package spanset

import (
	"context"
	"fmt"
	"strings"

	"github.com/cockroachdb/cockroach/pkg/keys"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/pkg/errors"
)

// SpanAccess records the intended mode of access in a SpanSet.
type SpanAccess int

// Constants for SpanAccess. Higher-valued accesses imply lower-level ones.
const (
	SpanReadOnly SpanAccess = iota
	SpanReadWrite
	NumSpanAccess
)

// String returns a string representation of the SpanAccess.
func (a SpanAccess) String() string {
	switch a {
	case SpanReadOnly:
		return "read"
	case SpanReadWrite:
		return "write"
	default:
		panic("unreachable")
	}
}

// SpanScope divides access types into local and global keys.
type SpanScope int

// Constants for span scopes.
const (
	SpanGlobal SpanScope = iota
	SpanLocal
	NumSpanScope
)

// String returns a string representation of the SpanScope.
func (a SpanScope) String() string {
	switch a {
	case SpanGlobal:
		return "global"
	case SpanLocal:
		return "local"
	default:
		panic("unreachable")
	}
}

// Span is used to represent a keyspan accessed by a request at a given
// timestamp. A zero timestamp indicates it's a non-MVCC access.
type Span struct {
	roachpb.Span
	Timestamp hlc.Timestamp
}

// SpanSet tracks the set of key spans touched by a command, broken into MVCC
// and non-MVCC accesses. The set is divided into subsets for access type
// (read-only or read/write) and key scope (local or global; used to facilitate
// use by the separate local and global latches).
// The Span slice for a particular access and scope contains non-ovelapping
// spans in increasing key order.
type SpanSet struct {
	spans [NumSpanAccess][NumSpanScope][]Span
}

// String prints a string representation of the SpanSet.
func (s *SpanSet) String() string {
	var buf strings.Builder
	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
		for ss := SpanScope(0); ss < NumSpanScope; ss++ {
			for _, cur := range s.GetSpans(sa, ss) {
				fmt.Fprintf(&buf, "%s %s: %s at %s\n",
					sa, ss, cur.Span.String(), cur.Timestamp.String())
			}
		}
	}
	return buf.String()
}

// Len returns the total number of spans tracked across all accesses and scopes.
func (s *SpanSet) Len() int {
	var count int
	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
		for ss := SpanScope(0); ss < NumSpanScope; ss++ {
			count += len(s.GetSpans(sa, ss))
		}
	}
	return count
}

// Reserve space for N additional spans.
func (s *SpanSet) Reserve(access SpanAccess, scope SpanScope, n int) {
	existing := s.spans[access][scope]
	s.spans[access][scope] = make([]Span, len(existing), n+cap(existing))
	copy(s.spans[access][scope], existing)
}

// AddNonMVCC adds a non-MVCC span to the span set. This should typically
// local keys.
func (s *SpanSet) AddNonMVCC(access SpanAccess, span roachpb.Span) {
	s.AddMVCC(access, span, hlc.Timestamp{})
}

// AddMVCC adds an MVCC span to the span set to be accessed at the given
// timestamp. This should typically be used for MVCC keys, user keys for e.g.
func (s *SpanSet) AddMVCC(access SpanAccess, span roachpb.Span, timestamp hlc.Timestamp) {
	scope := SpanGlobal
	if keys.IsLocal(span.Key) {
		scope = SpanLocal
	}

	s.spans[access][scope] = append(s.spans[access][scope], Span{Span: span, Timestamp: timestamp})
}

// SortAndDedup sorts the spans in the SpanSet and removes any duplicates.
func (s *SpanSet) SortAndDedup() {
	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
		for ss := SpanScope(0); ss < NumSpanScope; ss++ {
			s.spans[sa][ss], _ /* distinct */ = mergeSpans(s.spans[sa][ss])
		}
	}
}

// GetSpans returns a slice of spans with the given parameters.
func (s *SpanSet) GetSpans(access SpanAccess, scope SpanScope) []Span {
	return s.spans[access][scope]
}

// BoundarySpan returns a span containing all the spans with the given params.
func (s *SpanSet) BoundarySpan(scope SpanScope) roachpb.Span {
	var boundary roachpb.Span
	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
		for _, cur := range s.GetSpans(sa, scope) {
			if !boundary.Valid() {
				boundary = cur.Span
				continue
			}
			boundary = boundary.Combine(cur.Span)
		}
	}
	return boundary
}

// AssertAllowed calls CheckAllowed and fatals if the access is not allowed.
// Timestamps associated with the spans in the spanset are not considered,
// only the span boundaries are checked.
func (s *SpanSet) AssertAllowed(access SpanAccess, span roachpb.Span) {
	if err := s.CheckAllowed(access, span); err != nil {
		log.Fatal(context.TODO(), err)
	}
}

// CheckAllowed returns an error if the access is not allowed over the given
// keyspan. Timestamps associated with the spans in the spanset are not
// considered, only the span boundaries are checked.
//
// TODO(irfansharif): This does not currently work for spans that straddle
// across multiple added spans. Specifically a spanset with spans [a-c) and
// [b-d) added under read only and read write access modes respectively would
// fail at checking if read only access over the span [a-d) was requested. This
// is also a problem if the added spans were read only and the spanset wasn't
// already SortAndDedup-ed.
func (s *SpanSet) CheckAllowed(access SpanAccess, span roachpb.Span) error {
	return s.checkAllowed(access, span, false /* spanKeyExclusive */)
}

// See CheckAllowed(). The reversed arguments makes the lower bound exclusive
// and the upper bound inclusive, i.e. [a,b) will be considered (a,b].
func (s *SpanSet) checkAllowed(access SpanAccess, span roachpb.Span, reversed bool) error {
	scope := SpanGlobal
	if keys.IsLocal(span.Key) {
		scope = SpanLocal
	}

	for ac := access; ac < NumSpanAccess; ac++ {
		for _, cur := range s.spans[ac][scope] {
			if cur.Contains(span) &&
				(!reversed || cur.EndKey != nil && !cur.Key.Equal(span.Key)) ||
				reversed && cur.EndKey.Equal(span.Key) {
				return nil
			}
		}
	}

	return errors.Errorf("cannot %s undeclared span %s\ndeclared:\n%s", access, span, s)
}

// CheckAllowedAt returns an error if the access is not allowed at over the given keyspan
// at the given timestamp.
func (s *SpanSet) CheckAllowedAt(
	access SpanAccess, span roachpb.Span, timestamp hlc.Timestamp,
) error {
	return s.checkAllowedAt(access, span, timestamp, false /* inclusiveEnd */)
}

// See CheckAllowedAt. The reversed arguments makes the lower bound exclusive
// and the upper bound inclusive, i.e. [a,b) will be considered (a,b].
func (s *SpanSet) checkAllowedAt(
	access SpanAccess, span roachpb.Span, timestamp hlc.Timestamp, reversed bool,
) error {
	scope := SpanGlobal
	if keys.IsLocal(span.Key) {
		scope = SpanLocal
	}

	for ac := access; ac < NumSpanAccess; ac++ {
		for _, cur := range s.spans[ac][scope] {
			if (cur.Contains(span) &&
				(!reversed || (cur.EndKey != nil && !cur.Key.Equal(span.Key)))) ||
				(reversed && cur.EndKey.Equal(span.Key)) {
				if cur.Timestamp.IsEmpty() {
					// When the span is acquired as non-MVCC (i.e. with an empty
					// timestamp), it's equivalent to a read/write mutex where we don't
					// consider access timestamps.
					return nil
				}

				if access == SpanReadWrite {
					// Writes under a write span with an associated timestamp at that
					// specific timestamp.
					if timestamp == cur.Timestamp {
						return nil
					}
				} else {
					// Read spans acquired at a specific timestamp only allow reads at
					// that timestamp and below. Non-MVCC access is not allowed.
					if !timestamp.IsEmpty() && (timestamp.Less(cur.Timestamp) || timestamp == cur.Timestamp) {
						return nil
					}
				}
			}
		}
	}

	return errors.Errorf("cannot %s undeclared span %s at %s\ndeclared:\n%s",
		access, span, timestamp.String(), s)
}

// Validate returns an error if any spans that have been added to the set
// are invalid.
func (s *SpanSet) Validate() error {
	for sa := SpanAccess(0); sa < NumSpanAccess; sa++ {
		for ss := SpanScope(0); ss < NumSpanScope; ss++ {
			for _, cur := range s.GetSpans(sa, ss) {
				if len(cur.EndKey) > 0 && cur.Key.Compare(cur.EndKey) >= 0 {
					return errors.Errorf("inverted span %s %s", cur.Key, cur.EndKey)
				}
			}
		}
	}

	return nil
}