steve/pkg/counts/types.go

package counts

import (
	"context"
	"net/http"
	"sync"
	"time"

	"github.com/rancher/norman/pkg/store/empty"
	"github.com/rancher/norman/pkg/types"
	"golang.org/x/sync/errgroup"
)

var (
	ignore = map[string]bool{
		"count":  true,
		"schema": true,
	}
)

func Register(schemas *types.Schemas) {
	schemas.MustImportAndCustomize(Count{}, func(schema *types.Schema) {
		schema.CollectionMethods = []string{http.MethodGet}
		schema.ResourceMethods = []string{}
		schema.Store = &Store{}
	})
}

type Count struct {
	Counts map[string]ItemCount `json:"counts,omitempty"`
}

type ItemCount struct {
	Count    int    `json:"count,omitempty"`
	Revision string `json:"revision,omitempty"`
}

type Store struct {
	empty.Store
}

func (s *Store) ByID(apiOp *types.APIRequest, schema *types.Schema, id string) (types.APIObject, error) {
	c, err := s.getCount(apiOp, 750*time.Millisecond)
	return types.ToAPI(c), err
}

func (s *Store) List(apiOp *types.APIRequest, schema *types.Schema, opt *types.QueryOptions) (types.APIObject, error) {
	c, err := s.getCount(apiOp, 750*time.Millisecond)
	return types.ToAPI([]interface{}{c}), err
}

func (s *Store) Watch(apiOp *types.APIRequest, schema *types.Schema, w types.WatchRequest) (chan types.APIEvent, error) {
	return nil, nil
}

func (s *Store) getCount(apiOp *types.APIRequest, timeout time.Duration) (Count, error) {
	var countLock sync.Mutex
	counts := map[string]ItemCount{}

	errCtx, cancel := context.WithTimeout(apiOp.Context(), timeout)
	eg, errCtx := errgroup.WithContext(errCtx)
	defer cancel()

	for _, schema := range apiOp.Schemas.Schemas() {
		if ignore[schema.ID] {
			continue
		}

		if schema.Store == nil {
			continue
		}

		if apiOp.AccessControl.CanList(apiOp, schema) != nil {
			continue
		}

		current := schema
		eg.Go(func() error {
			list, err := current.Store.List(apiOp, current, nil)
			if err != nil {
				return err
			}
			if list.IsNil() {
				return nil
			}

			countLock.Lock()
			counts[current.ID] = ItemCount{
				Count:    len(list.List()),
				Revision: list.ListRevision,
			}
			countLock.Unlock()

			return nil
		})
	}

	if err := eg.Wait(); err != nil && err != context.Canceled {
		return Count{}, err
	}

	// in the case of cancellation go routines could still be running so we copy the map
	// to avoid returning a map that might get modified
	countLock.Lock()
	result := Count{
		Counts: map[string]ItemCount{},
	}
	for k, v := range counts {
		result.Counts[k] = v
	}
	countLock.Unlock()

	return result, nil
}