mutation_event.go

// Copyright 2023 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 event

import (
	"context"
	"fmt"

	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/util/log"
)

// SetSpanConfigEvent represents a mutation event responsible for updating the
// span config for ranges represented by the span.
type SetSpanConfigEvent struct {
	Span   roachpb.Span
	Config roachpb.SpanConfig
}

// AddNodeEvent represents a mutation event responsible for adding a node with
// its store count and locality specified by NumStores and LocalityString.
type AddNodeEvent struct {
	NumStores      int
	LocalityString string
}

// SetNodeLivenessEvent represents a mutation event responsible for setting
// liveness status of a node identified by the NodeID to the specified
// LivenessStatus.
type SetNodeLivenessEvent struct {
	NodeId         state.NodeID
	LivenessStatus livenesspb.NodeLivenessStatus
}

// SetCapacityOverrideEvent represents a mutation event responsible for updating
// the capacity for a store identified by StoreID.
type SetCapacityOverrideEvent struct {
	StoreID          state.StoreID
	CapacityOverride state.CapacityOverride
}

var _ Event = &SetSpanConfigEvent{}
var _ Event = &AddNodeEvent{}
var _ Event = &SetNodeLivenessEvent{}
var _ Event = &SetCapacityOverrideEvent{}

func (se SetSpanConfigEvent) Func() EventFunc {
	var fn MutationFunc = func(ctx context.Context, s state.State) {
		s.SetSpanConfig(se.Span, se.Config)
	}
	return fn
}

func (se SetSpanConfigEvent) String() string {
	return ""
}

func (ae AddNodeEvent) Func() EventFunc {
	var fn MutationFunc = func(ctx context.Context, s state.State) {
		node := s.AddNode()
		if ae.LocalityString != "" {
			var locality roachpb.Locality
			if err := locality.Set(ae.LocalityString); err != nil {
				panic(fmt.Sprintf("unable to set node locality %s", err.Error()))
			}
			s.SetNodeLocality(node.NodeID(), locality)
		}
		for i := 0; i < ae.NumStores; i++ {
			if _, ok := s.AddStore(node.NodeID()); !ok {
				panic(fmt.Sprintf("adding store to node=%d failed", node))
			}
		}
	}
	return fn
}

func (ae AddNodeEvent) String() string {
	return ""
}

func (sne SetNodeLivenessEvent) Func() EventFunc {
	var fn MutationFunc = func(ctx context.Context, s state.State) {
		s.SetNodeLiveness(
			sne.NodeId,
			sne.LivenessStatus,
		)
	}
	return fn
}

func (sne SetNodeLivenessEvent) String() string {
	return "'"
}

func (sce SetCapacityOverrideEvent) Func() EventFunc {
	var fn MutationFunc = func(ctx context.Context, s state.State) {
		log.Infof(ctx, "setting capacity override %v", sce.CapacityOverride)
		s.SetCapacityOverride(sce.StoreID, sce.CapacityOverride)
	}
	return fn
}

func (sce SetCapacityOverrideEvent) String() string {
	return "'"
}