gopkg.in/ubuntu-core/snappy.v0@v0.0.0-20210902073436-25a8614f10a6/overlord/servicestate/internal/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 internal
    21  
    22  import (
    23  	"fmt"
    24  	"sort"
    25  
    26  	"github.com/snapcore/snapd/gadget/quantity"
    27  	"github.com/snapcore/snapd/overlord/state"
    28  	"github.com/snapcore/snapd/snap/quota"
    29  )
    30  
    31  // AllQuotas returns all currently tracked quota groups in the state. They are
    32  // validated for consistency using ResolveCrossReferences before being returned.
    33  func AllQuotas(st *state.State) (map[string]*quota.Group, error) {
    34  	var quotas map[string]*quota.Group
    35  	if err := st.Get("quotas", &quotas); err != nil {
    36  		if err != state.ErrNoState {
    37  			return nil, err
    38  		}
    39  		// otherwise there are no quotas so just return nil
    40  		return nil, nil
    41  	}
    42  
    43  	// quota groups are not serialized with all the necessary tracking
    44  	// information in the objects, so we need to thread some things around
    45  	if err := quota.ResolveCrossReferences(quotas); err != nil {
    46  		return nil, err
    47  	}
    48  
    49  	// quotas has now been properly initialized with unexported cross-references
    50  	return quotas, nil
    51  }
    52  
    53  // PatchQuotas will update the state quota group map with the provided quota
    54  // groups. It returns the full set of all quota groups after a successful
    55  // update for convenience. The groups provided will replace group states if
    56  // present or be added on top of the current set of quota groups in the state,
    57  // and verified for consistency before committed to state. When adding
    58  // sub-groups, both the parent and the sub-group must be added at once since the
    59  // sub-group needs to reference the parent group and vice versa to be fully
    60  // consistent.
    61  func PatchQuotas(st *state.State, grps ...*quota.Group) (map[string]*quota.Group, error) {
    62  	// get the current set of quotas
    63  	allGrps, err := AllQuotas(st)
    64  	if err != nil {
    65  		// AllQuotas() can't return ErrNoState, in that case it just returns a
    66  		// nil map, which we handle below
    67  		return nil, err
    68  	}
    69  	if allGrps == nil {
    70  		allGrps = make(map[string]*quota.Group)
    71  	}
    72  
    73  	// handle trivial case here to prevent panics below
    74  	if len(grps) == 0 {
    75  		return allGrps, nil
    76  	}
    77  
    78  	sort.SliceStable(grps, func(i, j int) bool {
    79  		return grps[i].Name < grps[j].Name
    80  	})
    81  
    82  	// add to the temporary state map
    83  	for _, grp := range grps {
    84  		allGrps[grp.Name] = grp
    85  	}
    86  
    87  	// make sure the full set is still resolved before saving it - this prevents
    88  	// easy errors like trying to add a sub-group quota without updating the
    89  	// parent with references to the sub-group, for cases like those, all
    90  	// related groups must be updated at the same time in one operation to
    91  	// prevent having inconsistent quota groups in state.json
    92  	if err := quota.ResolveCrossReferences(allGrps); err != nil {
    93  		// make a nice error message for this case
    94  		updated := ""
    95  		for _, grp := range grps[:len(grps)-1] {
    96  			updated += fmt.Sprintf("%q, ", grp.Name)
    97  		}
    98  		updated += fmt.Sprintf("%q", grps[len(grps)-1].Name)
    99  		plural := ""
   100  		if len(grps) > 1 {
   101  			plural = "s"
   102  		}
   103  		return nil, fmt.Errorf("cannot update quota%s %s: %v", plural, updated, err)
   104  	}
   105  
   106  	st.Set("quotas", allGrps)
   107  	return allGrps, nil
   108  }
   109  
   110  // CreateQuotaInState creates a quota group with the given paremeters
   111  // in the state.  It takes the current map of all quota groups.
   112  func CreateQuotaInState(st *state.State, quotaName string, parentGrp *quota.Group, snaps []string, memoryLimit quantity.Size, allGrps map[string]*quota.Group) (*quota.Group, map[string]*quota.Group, error) {
   113  	// make sure that the parent group exists if we are creating a sub-group
   114  	var grp *quota.Group
   115  	var err error
   116  	updatedGrps := []*quota.Group{}
   117  	if parentGrp != nil {
   118  		grp, err = parentGrp.NewSubGroup(quotaName, memoryLimit)
   119  		if err != nil {
   120  			return nil, nil, err
   121  		}
   122  
   123  		updatedGrps = append(updatedGrps, parentGrp)
   124  	} else {
   125  		// make a new group
   126  		grp, err = quota.NewGroup(quotaName, memoryLimit)
   127  		if err != nil {
   128  			return nil, nil, err
   129  		}
   130  	}
   131  	updatedGrps = append(updatedGrps, grp)
   132  
   133  	// put the snaps in the group
   134  	grp.Snaps = snaps
   135  	// update the modified groups in state
   136  	newAllGrps, err := PatchQuotas(st, updatedGrps...)
   137  	if err != nil {
   138  		return nil, nil, err
   139  	}
   140  
   141  	return grp, newAllGrps, nil
   142  }