services/mounttable/btmtd/internal/counters.go

// Copyright 2016 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package internal

import (
	"bytes"
	"encoding/binary"

	"google.golang.org/cloud/bigtable"

	"v.io/v23/context"
	"v.io/v23/verror"
)

var (
	errNodeLimitExceeded   = verror.Register(pkgPath+".errNodeLimitExceeded", verror.NoRetry, "{1:}{2:} node limit per user ({3}) exceeded{:_}")
	errServerLimitExceeded = verror.Register(pkgPath+".errServerLimitExceeded", verror.NoRetry, "{1:}{2:} mounted server limit per user ({3}) exceeded{:_}")
)

func incrementCounter(ctx *context.T, bt *BigTable, name string, delta int64) (int64, error) {
	bctx, cancel := btctx(ctx)
	defer cancel()

	m := bigtable.NewReadModifyWrite()
	m.Increment(metadataFamily, "c", delta)
	row, err := bt.counterTbl.ApplyReadModifyWrite(bctx, name, m)
	if err != nil {
		return 0, err
	}
	return decodeCounterValue(ctx, row)
}

func decodeCounterValue(ctx *context.T, row bigtable.Row) (c int64, err error) {
	if len(row[metadataFamily]) != 1 {
		return 0, verror.NewErrInternal(ctx)
	}
	b := row[metadataFamily][0].Value
	err = binary.Read(bytes.NewReader(b), binary.BigEndian, &c)
	return
}

func incrementWithLimit(ctx *context.T, bt *BigTable, name string, delta, limit int64, overLimitError verror.IDAction) error {
	if delta == 0 {
		return nil
	}
	c, err := incrementCounter(ctx, bt, name, delta)
	if err != nil {
		return err
	}
	if delta > 0 && limit > 0 && c > limit {
		if _, err := incrementCounter(ctx, bt, name, -delta); err != nil {
			ctx.Errorf("incrementCounter(%q, %v) failed: %v", name, -delta, err)
		}
		return verror.New(overLimitError, ctx, limit)
	}
	return nil
}

func incrementCreatorNodeCount(ctx *context.T, bt *BigTable, creator string, delta, limit int64) error {
	return incrementWithLimit(ctx, bt, "num-nodes-per-user:"+creator, delta, limit, errNodeLimitExceeded)
}

func incrementCreatorServerCount(ctx *context.T, bt *BigTable, creator string, delta, limit int64) error {
	return incrementWithLimit(ctx, bt, "num-servers-per-user:"+creator, delta, limit, errServerLimitExceeded)
}

func recalculateCounters(ctx *context.T, bt *BigTable) error {
	bctx, cancel := btctx(ctx)
	defer cancel()

	// Delete all the counters.
	if err := bt.counterTbl.ReadRows(bctx, bigtable.InfiniteRange(""),
		func(row bigtable.Row) bool {
			mut := bigtable.NewMutation()
			mut.DeleteRow()
			if err := bt.counterTbl.Apply(ctx, row.Key(), mut); err != nil {
				ctx.Errorf("apply delete row (%q) failed: %v", row.Key(), err)
				return false
			}
			return true
		},
	); err != nil {
		return err
	}

	// Re-create all the counters.
	return bt.nodeTbl.ReadRows(bctx, bigtable.InfiniteRange(""),
		func(row bigtable.Row) bool {
			n := nodeFromRow(ctx, bt, row, clock)
			if err := incrementCreatorNodeCount(ctx, bt, n.creator, 1, 0); err != nil {
				ctx.Errorf("incrementCreatorNodeCount(%q) failed: %v", n.name, err)
				return false
			}
			if err := incrementCreatorServerCount(ctx, bt, n.creator, int64(len(n.servers)+len(n.expiredServers)), 0); err != nil {
				ctx.Errorf("incrementCreatorServerCount(%q) failed: %v", n.name, err)
				return false
			}
			return true
		},
		bigtable.RowFilter(bigtable.LatestNFilter(1)),
	)
}