vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/rules/map.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package rules
    18  
    19  import (
    20  	"encoding/json"
    21  	"errors"
    22  	"sync"
    23  
    24  	"vitess.io/vitess/go/vt/log"
    25  	"vitess.io/vitess/go/vt/vttablet/tabletserver/planbuilder"
    26  )
    27  
    28  // Map is the maintainer of Rules from multiple sources
    29  type Map struct {
    30  	// mutex to protect following queryRulesMap
    31  	mu sync.Mutex
    32  	// queryRulesMap maps the names of different query rule sources to the actual Rules structure
    33  	queryRulesMap map[string]*Rules
    34  }
    35  
    36  // NewMap returns an empty Map object.
    37  func NewMap() *Map {
    38  	qri := &Map{
    39  		queryRulesMap: map[string]*Rules{},
    40  	}
    41  	return qri
    42  }
    43  
    44  // RegisterSource registers a query rule source name with Map.
    45  func (qri *Map) RegisterSource(ruleSource string) {
    46  	qri.mu.Lock()
    47  	defer qri.mu.Unlock()
    48  	if _, existed := qri.queryRulesMap[ruleSource]; existed {
    49  		log.Errorf("Query rule source " + ruleSource + " has been registered")
    50  		panic("Query rule source " + ruleSource + " has been registered")
    51  	}
    52  	qri.queryRulesMap[ruleSource] = New()
    53  }
    54  
    55  // UnRegisterSource removes a registered query rule source name.
    56  func (qri *Map) UnRegisterSource(ruleSource string) {
    57  	qri.mu.Lock()
    58  	defer qri.mu.Unlock()
    59  	delete(qri.queryRulesMap, ruleSource)
    60  }
    61  
    62  // SetRules takes an external Rules structure and overwrite one of the
    63  // internal Rules as designated by ruleSource parameter.
    64  func (qri *Map) SetRules(ruleSource string, newRules *Rules) error {
    65  	if newRules == nil {
    66  		newRules = New()
    67  	}
    68  	qri.mu.Lock()
    69  	defer qri.mu.Unlock()
    70  	if _, ok := qri.queryRulesMap[ruleSource]; ok {
    71  		qri.queryRulesMap[ruleSource] = newRules.Copy()
    72  		return nil
    73  	}
    74  	return errors.New("Rule source identifier " + ruleSource + " is not valid")
    75  }
    76  
    77  // Get returns the corresponding Rules as designated by ruleSource parameter.
    78  func (qri *Map) Get(ruleSource string) (*Rules, error) {
    79  	qri.mu.Lock()
    80  	defer qri.mu.Unlock()
    81  	if ruleset, ok := qri.queryRulesMap[ruleSource]; ok {
    82  		return ruleset.Copy(), nil
    83  	}
    84  	return New(), errors.New("Rule source identifier " + ruleSource + " is not valid")
    85  }
    86  
    87  // FilterByPlan creates a new Rules by prefiltering on all query rules that are contained in internal
    88  // Rules structures, in other words, query rules from all predefined sources will be applied.
    89  func (qri *Map) FilterByPlan(query string, planid planbuilder.PlanType, tableNames ...string) (newqrs *Rules) {
    90  	qri.mu.Lock()
    91  	defer qri.mu.Unlock()
    92  	newqrs = New()
    93  	for _, rules := range qri.queryRulesMap {
    94  		newqrs.Append(rules.FilterByPlan(query, planid, tableNames...))
    95  	}
    96  	return newqrs
    97  }
    98  
    99  // MarshalJSON marshals to JSON.
   100  func (qri *Map) MarshalJSON() ([]byte, error) {
   101  	qri.mu.Lock()
   102  	defer qri.mu.Unlock()
   103  	return json.Marshal(qri.queryRulesMap)
   104  }