github.com/freetocompute/snapd@v0.0.0-20210618182524-2fb355d72fd9/overlord/servicestate/quotas.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2021 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package servicestate
    21  
    22  import (
    23  	"errors"
    24  	"fmt"
    25  	"sort"
    26  
    27  	"github.com/snapcore/snapd/overlord/state"
    28  	"github.com/snapcore/snapd/snap/quota"
    29  )
    30  
    31  var ErrQuotaNotFound = errors.New("quota not found")
    32  
    33  // AllQuotas returns all currently tracked quota groups in the state. They are
    34  // validated for consistency using ResolveCrossReferences before being returned.
    35  func AllQuotas(st *state.State) (map[string]*quota.Group, error) {
    36  	var quotas map[string]*quota.Group
    37  	if err := st.Get("quotas", &quotas); err != nil {
    38  		if err != state.ErrNoState {
    39  			return nil, err
    40  		}
    41  		// otherwise there are no quotas so just return nil
    42  		return nil, nil
    43  	}
    44  
    45  	// quota groups are not serialized with all the necessary tracking
    46  	// information in the objects, so we need to thread some things around
    47  	if err := quota.ResolveCrossReferences(quotas); err != nil {
    48  		return nil, err
    49  	}
    50  
    51  	// quotas has now been properly initialized with unexported cross-references
    52  	return quotas, nil
    53  }
    54  
    55  // GetQuota returns an individual quota group by name.
    56  func GetQuota(st *state.State, name string) (*quota.Group, error) {
    57  	allGrps, err := AllQuotas(st)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  
    62  	group, ok := allGrps[name]
    63  	if !ok {
    64  		return nil, ErrQuotaNotFound
    65  	}
    66  
    67  	return group, nil
    68  }
    69  
    70  // patchQuotas will update the state quota group map with the provided quota
    71  // groups. It returns the full set of all quota groups after a successful
    72  // update for convenience. The groups provided will replace group states if
    73  // present or be added on top of the current set of quota groups in the state,
    74  // and verified for consistency before committed to state. When adding
    75  // sub-groups, both the parent and the sub-group must be added at once since the
    76  // sub-group needs to reference the parent group and vice versa to be fully
    77  // consistent.
    78  func patchQuotas(st *state.State, grps ...*quota.Group) (map[string]*quota.Group, error) {
    79  	// get the current set of quotas
    80  	allGrps, err := AllQuotas(st)
    81  	if err != nil {
    82  		// AllQuotas() can't return ErrNoState, in that case it just returns a
    83  		// nil map, which we handle below
    84  		return nil, err
    85  	}
    86  	if allGrps == nil {
    87  		allGrps = make(map[string]*quota.Group)
    88  	}
    89  
    90  	// handle trivial case here to prevent panics below
    91  	if len(grps) == 0 {
    92  		return allGrps, nil
    93  	}
    94  
    95  	sort.SliceStable(grps, func(i, j int) bool {
    96  		return grps[i].Name < grps[j].Name
    97  	})
    98  
    99  	// add to the temporary state map
   100  	for _, grp := range grps {
   101  		allGrps[grp.Name] = grp
   102  	}
   103  
   104  	// make sure the full set is still resolved before saving it - this prevents
   105  	// easy errors like trying to add a sub-group quota without updating the
   106  	// parent with references to the sub-group, for cases like those, all
   107  	// related groups must be updated at the same time in one operation to
   108  	// prevent having inconsistent quota groups in state.json
   109  	if err := quota.ResolveCrossReferences(allGrps); err != nil {
   110  		// make a nice error message for this case
   111  		updated := ""
   112  		for _, grp := range grps[:len(grps)-1] {
   113  			updated += fmt.Sprintf("%q, ", grp.Name)
   114  		}
   115  		updated += fmt.Sprintf("%q", grps[len(grps)-1].Name)
   116  		plural := ""
   117  		if len(grps) > 1 {
   118  			plural = "s"
   119  		}
   120  		return nil, fmt.Errorf("cannot update quota%s %s: %v", plural, updated, err)
   121  	}
   122  
   123  	st.Set("quotas", allGrps)
   124  	return allGrps, nil
   125  }