<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/litl/galaxy/registry"
	"github.com/litl/galaxy/utils"
	log "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	serviceRegistry *registry.ServiceRegistry
	router          *hostroute.HostRouter
	env             string
	version         bool
	buildVersion    string
	listenAddr      string
	redisHost       string
)

type RequestLogger struct{}

func (r *RequestLogger) ObserveRequest(req request.Request) {
	log.Infof("%s %s", req.GetHttpRequest().Host, req)
}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {}

var balancers map[string]*roundrobin.RoundRobin

func addBackends(registrations []registry.ServiceRegistration) map[string][]string {

	liveVhosts := make(map[string][]string)
	for _, r := range registrations {

		serviceConfig, err := serviceRegistry.GetServiceConfig(r.Name)
		if err != nil {
			log.Errorf("Error: %s", err)
			continue
		}
		vhosts := serviceConfig.Env()["VIRTUAL_HOST"]

		for _, vhost := range strings.Split(vhosts, ",") {

			if vhost == "" {
				continue
			}

			addr := fmt.Sprint(r.ExternalIP, ":", r.ExternalPort)
			url := "http://" + addr
			liveVhosts[vhost] = append(liveVhosts[vhost], url)
			balancer := balancers[vhost]

			if balancer == nil {
				// Create a round robin load balancer with some endpoints
				balancer, err = roundrobin.NewRoundRobin()
				if err != nil {
					log.Errorf("Error: %s", err)
					continue
				}

				// Create a http location with the load balancer we've just added
				loc, err := httploc.NewLocationWithOptions(r.Name, balancer,
					httploc.Options{
						TrustForwardHeader: true,
					})
				if err != nil {
					log.Errorf("Error: %s", err)
					continue
				}
				loc.GetObserverChain().Add("logger", &RequestLogger{})

				router.SetRouter(vhost, &route.ConstRouter{Location: loc})
				log.Infof("Creating balancer for %s", vhost)
				balancers[vhost] = balancer
			}

			// Already registered?
			if balancer.FindEndpointByUrl(url) != nil {
				continue
			}
			endpoint := endpoint.MustParseUrl(url)
			log.Infof("Adding %s endpoint %s", vhost, endpoint.GetUrl())
			err := balancer.AddEndpoint(endpoint)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
	}
	return liveVhosts
}

func removeBackends(liveVhosts map[string][]string) {
	// Remove backends that are no longer registered
	for k, _ := range balancers {

		if k == "" {
			continue
		}
		balancer := balancers[k]
		endpoints := balancer.GetEndpoints()
		for _, endpoint := range endpoints {
			vhosts := liveVhosts[k]
			exists := false
			for _, url := range vhosts {
				if endpoint.GetUrl().String() == url {
					exists = true
					break
				}
			}
			if !exists {
				log.Infof("Removing %s endpoint %s", k, endpoint.GetUrl())
				balancer.RemoveEndpoint(endpoint)
			}
		}

		endpoints = balancer.GetEndpoints()
		if len(endpoints) == 0 {
			log.Infof("Removing balancer for %s", k)
			delete(balancers, k)
			router.RemoveRouter(k)
			continue
		}

	}
}

func updateRoutes() {
	for {
		registrations, err := serviceRegistry.ListRegistrations()
		if err != nil {
			log.Errorf("%s", err)
			return
		}

		liveVhosts := addBackends(registrations)
		removeBackends(liveVhosts)
		time.Sleep(10 * time.Second)
	}
}

func adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(balancers) == 0 {
		w.WriteHeader(503)
		return
	}
	for k, _ := range balancers {
		balancer := balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warningf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := balancers[host]; !exists {
			adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func main() {

	flag.StringVar(&listenAddr, "http", "127.0.0.1:8080", "http server address")
	flag.StringVar(&env, "env", utils.GetEnv("GALAXY_ENV", ""), "Environment namespace")
	flag.StringVar(&redisHost, "redis", utils.GetEnv("GALAXY_REDIS_HOST", "127.0.0.1:6379"), "redis host")
	flag.BoolVar(&version, "v", false, "display version info")

	flag.Parse()

	if version {
		fmt.Println(buildVersion)
		return
	}

	if env == "" {
		fmt.Println("Need an env")
		flag.PrintDefaults()
		os.Exit(1)
	}

	balancers = make(map[string]*roundrobin.RoundRobin)
	serviceRegistry = registry.NewServiceRegistry(
		env,
		"",
		"",
		30,
		"",
	)

	serviceRegistry.Connect(redisHost)

	outputBuffer := &utils.OutputBuffer{}
	serviceRegistry.OutputBuffer = outputBuffer

	// init the logging package
	log.Init([]*log.LogConfig{
		&log.LogConfig{Name: "console"},
	})

	log.Infof("Starting quasar for env %s", env)
	log.Infof("Listening at %s", listenAddr)

	router = hostroute.NewHostRouter()
	go updateRoutes()

	proxy, err := vulcan.NewProxy(router)
	if err != nil {
		log.Fatalf("Error: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}
<commit_msg>Don't modify map we're iterating over in the loop<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/litl/galaxy/registry"
	"github.com/litl/galaxy/utils"
	log "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	serviceRegistry *registry.ServiceRegistry
	router          *hostroute.HostRouter
	env             string
	version         bool
	buildVersion    string
	listenAddr      string
	redisHost       string
)

type RequestLogger struct{}

func (r *RequestLogger) ObserveRequest(req request.Request) {
	log.Infof("%s %s", req.GetHttpRequest().Host, req)
}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {}

var balancers map[string]*roundrobin.RoundRobin

func addBackends(registrations []registry.ServiceRegistration) map[string][]string {

	liveVhosts := make(map[string][]string)
	for _, r := range registrations {

		serviceConfig, err := serviceRegistry.GetServiceConfig(r.Name)
		if err != nil {
			log.Errorf("Error: %s", err)
			continue
		}
		vhosts := serviceConfig.Env()["VIRTUAL_HOST"]

		for _, vhost := range strings.Split(vhosts, ",") {

			if vhost == "" {
				continue
			}

			addr := fmt.Sprint(r.ExternalIP, ":", r.ExternalPort)
			url := "http://" + addr
			liveVhosts[vhost] = append(liveVhosts[vhost], url)
			balancer := balancers[vhost]

			if balancer == nil {
				// Create a round robin load balancer with some endpoints
				balancer, err = roundrobin.NewRoundRobin()
				if err != nil {
					log.Errorf("Error: %s", err)
					continue
				}

				// Create a http location with the load balancer we've just added
				loc, err := httploc.NewLocationWithOptions(r.Name, balancer,
					httploc.Options{
						TrustForwardHeader: true,
					})
				if err != nil {
					log.Errorf("Error: %s", err)
					continue
				}
				loc.GetObserverChain().Add("logger", &RequestLogger{})

				router.SetRouter(vhost, &route.ConstRouter{Location: loc})
				log.Infof("Creating balancer for %s", vhost)
				balancers[vhost] = balancer
			}

			// Already registered?
			if balancer.FindEndpointByUrl(url) != nil {
				continue
			}
			endpoint := endpoint.MustParseUrl(url)
			log.Infof("Adding %s endpoint %s", vhost, endpoint.GetUrl())
			err := balancer.AddEndpoint(endpoint)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
	}
	return liveVhosts
}

func removeBackends(liveVhosts map[string][]string) {
	// Remove backends that are no longer registered

	remove := []string{}
	for k, _ := range balancers {

		if k == "" {
			continue
		}
		balancer := balancers[k]
		endpoints := balancer.GetEndpoints()
		for _, endpoint := range endpoints {
			vhosts := liveVhosts[k]
			exists := false
			for _, url := range vhosts {
				if endpoint.GetUrl().String() == url {
					exists = true
					break
				}
			}
			if !exists {
				log.Infof("Removing %s endpoint %s", k, endpoint.GetUrl())
				balancer.RemoveEndpoint(endpoint)
			}
		}

		endpoints = balancer.GetEndpoints()
		if len(endpoints) == 0 {
			remove = append(remove, k)
		}
	}

	for _, v := range remove {
		log.Infof("Removing balancer for %s", v)
		delete(balancers, v)
		router.RemoveRouter(v)
	}
}

func updateRoutes() {
	for {
		registrations, err := serviceRegistry.ListRegistrations()
		if err != nil {
			log.Errorf("%s", err)
			return
		}

		liveVhosts := addBackends(registrations)
		removeBackends(liveVhosts)
		time.Sleep(10 * time.Second)
	}
}

func adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(balancers) == 0 {
		w.WriteHeader(503)
		return
	}
	for k, _ := range balancers {
		balancer := balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warningf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := balancers[host]; !exists {
			adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func main() {

	flag.StringVar(&listenAddr, "http", "127.0.0.1:8080", "http server address")
	flag.StringVar(&env, "env", utils.GetEnv("GALAXY_ENV", ""), "Environment namespace")
	flag.StringVar(&redisHost, "redis", utils.GetEnv("GALAXY_REDIS_HOST", "127.0.0.1:6379"), "redis host")
	flag.BoolVar(&version, "v", false, "display version info")

	flag.Parse()

	if version {
		fmt.Println(buildVersion)
		return
	}

	if env == "" {
		fmt.Println("Need an env")
		flag.PrintDefaults()
		os.Exit(1)
	}

	balancers = make(map[string]*roundrobin.RoundRobin)
	serviceRegistry = registry.NewServiceRegistry(
		env,
		"",
		"",
		30,
		"",
	)

	serviceRegistry.Connect(redisHost)

	outputBuffer := &utils.OutputBuffer{}
	serviceRegistry.OutputBuffer = outputBuffer

	// init the logging package
	log.Init([]*log.LogConfig{
		&log.LogConfig{Name: "console"},
	})

	log.Infof("Starting quasar for env %s", env)
	log.Infof("Listening at %s", listenAddr)

	router = hostroute.NewHostRouter()
	go updateRoutes()

	proxy, err := vulcan.NewProxy(router)
	if err != nil {
		log.Fatalf("Error: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package quota

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

	"github.com/apigee/apigee-remote-service-golib/auth"
	"github.com/apigee/apigee-remote-service-golib/log"
	"github.com/apigee/apigee-remote-service-golib/product"
	"github.com/apigee/apigee-remote-service-golib/util"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

const (
	quotaPath             = "/quotas"
	defaultSyncRate       = time.Second
	defaultNumSyncWorkers = 10
	defaultRefreshAfter   = 1 * time.Minute
	defaultDeleteAfter    = 10 * time.Minute
	syncQueueSize         = 100
	resultCacheBufferSize = 30
)

/*
Dispatch basically works like this:
	every m.syncRate tick do
		if bucket needs to sync and isn't syncing
			m.bucketToSyncQueue <- bucket
	each of m.numSyncWorkers do:
		bucket <-m.bucketToSyncQueue
			bucket.sync() with backoff
*/

// A Manager tracks multiple Apigee quotas
type Manager interface {
	Start()
	Apply(auth *auth.Context, p *product.APIProduct, args Args) (*Result, error)
	Close()
}

type manager struct {
	baseURL            *url.URL
	close              chan bool
	client             *http.Client
	now                func() time.Time
	syncRate           time.Duration
	bucketsLock        sync.RWMutex
	buckets            map[string]*bucket // Map from ID -> bucket
	bucketToSyncQueue  chan *bucket
	numSyncWorkers     int
	syncWorkerWG       sync.WaitGroup
	dupCache           ResultCache
	bucketsSyncingLock sync.Mutex
	bucketsSyncing     map[*bucket]struct{}
	org                string
	env                string
	runningContext     context.Context
	cancelContext      context.CancelFunc
}

// NewManager constructs and starts a new Manager. Call Close when done.
func NewManager(options Options) (Manager, error) {
	if err := options.validate(); err != nil {
		return nil, err
	}
	m := newManager(options)
	m.Start()
	return m, nil
}

// newManager constructs a new Manager
func newManager(options Options) *manager {
	return &manager{
		close:             make(chan bool),
		client:            options.Client,
		now:               time.Now,
		syncRate:          defaultSyncRate,
		buckets:           map[string]*bucket{},
		bucketToSyncQueue: make(chan *bucket, syncQueueSize),
		baseURL:           options.BaseURL,
		numSyncWorkers:    defaultNumSyncWorkers,
		dupCache:          ResultCache{size: resultCacheBufferSize},
		bucketsSyncing:    map[*bucket]struct{}{},
		org:               options.Org,
		env:               options.Env,
	}
}

// Start starts the manager.
func (m *manager) Start() {
	log.Infof("starting quota manager")

	m.runningContext, m.cancelContext = context.WithCancel(context.Background())

	go m.bucketMaintenanceLoop()
	for i := 0; i < m.numSyncWorkers; i++ {
		go m.syncBucketDispatcher()
	}

	log.Infof("started quota manager with %d workers", m.numSyncWorkers)
}

// Close shuts down the manager.
func (m *manager) Close() {
	if m == nil {
		return
	}
	log.Infof("closing quota manager")
	m.cancelContext()
	m.runningContext = nil

	m.syncWorkerWG.Add(m.numSyncWorkers)
	m.close <- true
	close(m.bucketToSyncQueue)
	m.syncWorkerWG.Wait()
	log.Infof("closed quota manager")
}

func getQuotaID(auth *auth.Context, p *product.APIProduct) string {
	return fmt.Sprintf("%s-%s", auth.Application, p.Name)
}

// Apply a quota request to the local quota bucket and schedule for sync
func (m *manager) Apply(auth *auth.Context, p *product.APIProduct, args Args) (*Result, error) {

	if result := m.dupCache.Get(args.DeduplicationID); result != nil {
		return result, nil
	}

	quotaID := getQuotaID(auth, p)

	req := &Request{
		Identifier: quotaID,
		Interval:   p.QuotaIntervalInt,
		Allow:      p.QuotaLimitInt,
		TimeUnit:   p.QuotaTimeUnit,
	}

	// a new bucket is created if missing or if product is no longer compatible
	var result *Result
	var err error
	m.bucketsLock.RLock()
	b, ok := m.buckets[quotaID]
	m.bucketsLock.RUnlock()
	if !ok || !b.compatible(req) {
		m.bucketsLock.Lock()
		b, ok = m.buckets[quotaID]
		if !ok || !b.compatible(req) {
			promLabels := m.prometheusLabelsForQuota(quotaID)
			b = newBucket(*req, m, promLabels)
			m.buckets[quotaID] = b
			log.Debugf("new quota bucket: %s", quotaID)
		}
		m.bucketsLock.Unlock()
	}

	req.Weight = args.QuotaAmount
	result, err = b.apply(req)

	if result != nil && err == nil && args.DeduplicationID != "" {
		m.dupCache.Add(args.DeduplicationID, result)
	}

	return result, err
}

// loop to sync active buckets and delete old buckets
func (m *manager) bucketMaintenanceLoop() {
	t := time.NewTicker(m.syncRate)
	for {
		select {
		case <-t.C:
			var deleteIDs []string
			m.bucketsLock.RLock()
			for id, b := range m.buckets {
				if b.needToDelete() {
					deleteIDs = append(deleteIDs, id)

				} else if b.needToSync() {
					m.bucketsSyncingLock.Lock()
					if _, ok := m.bucketsSyncing[b]; !ok { // not already scheduled
						m.bucketsSyncing[b] = struct{}{}
						m.bucketToSyncQueue <- b
					}
					m.bucketsSyncingLock.Unlock()
				}
			}
			m.bucketsLock.RUnlock()

			if deleteIDs != nil {
				log.Debugf("deleting quota buckets: %v", deleteIDs)
				m.bucketsLock.Lock()
				for _, id := range deleteIDs {
					delete(m.buckets, id)
					labels := m.prometheusLabelsForQuota(id)
					prometheusBucketWindowExpires.Delete(labels)
					prometheusBucketChecked.Delete(labels)
					prometheusBucketSynced.Delete(labels)
					prometheusBucketValue.Delete(labels)
				}
				m.bucketsLock.Unlock()
			}

		case <-m.close:
			log.Debugf("closing quota sync loop")
			t.Stop()
			return
		}
	}
}

func (m *manager) prometheusLabelsForQuota(quotaID string) prometheus.Labels {
	return prometheus.Labels{"org": m.org, "env": m.env, "quota": quotaID}
}

// routine for dispatching work to sync a bucket with the server
func (m *manager) syncBucketDispatcher() {

	for bucket := range m.bucketToSyncQueue {
		looper := util.Looper{
			Backoff: util.NewExponentialBackoff(200*time.Millisecond, 30*time.Second, 2, true),
		}
		work := func(ctx context.Context) error {
			return bucket.sync()
		}
		errH := func(err error) error {
			log.Errorf("sync: %s", err)
			return nil
		}

		// run until success or canceled with backoff
		if err := looper.Run(m.runningContext, work, errH); err != nil {
			log.Errorf("looper run: %s", err)
		}

		m.bucketsSyncingLock.Lock()
		delete(m.bucketsSyncing, bucket)
		m.bucketsSyncingLock.Unlock()
	}

	m.syncWorkerWG.Done()
	log.Debugf("closing quota sync worker")
}

// Options allows us to specify options for how this auth manager will run
type Options struct {
	// Client is a configured HTTPClient
	Client *http.Client
	// BaseURL of the Apigee internal proxy
	BaseURL *url.URL
	// Org is organization
	Org string
	// Env is environment
	Env string
}

func (o *Options) validate() error {
	if o.Client == nil ||
		o.BaseURL == nil ||
		o.Org == "" ||
		o.Env == "" {
		return fmt.Errorf("all quota options are required")
	}
	return nil
}

var (
	prometheusBucketValue = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "value",
		Help:      "Current value of a quota",
	}, []string{"org", "env", "quota"})

	prometheusBucketChecked = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "checked",
		Help:      "Time quota was last checked",
	}, []string{"org", "env", "quota"})

	prometheusBucketSynced = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "synced",
		Help:      "Time quota was last synced",
	}, []string{"org", "env", "quota"})

	prometheusBucketWindowExpires = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "window_expires",
		Help:      "Time quota window will expire",
	}, []string{"org", "env", "quota"})
)
<commit_msg>address sync deadlock, increase queue size (#18)<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package quota

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

	"github.com/apigee/apigee-remote-service-golib/auth"
	"github.com/apigee/apigee-remote-service-golib/log"
	"github.com/apigee/apigee-remote-service-golib/product"
	"github.com/apigee/apigee-remote-service-golib/util"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

const (
	quotaPath             = "/quotas"
	defaultSyncRate       = time.Second
	defaultNumSyncWorkers = 10
	defaultRefreshAfter   = 1 * time.Minute
	defaultDeleteAfter    = 10 * time.Minute
	syncQueueSize         = 1000
	resultCacheBufferSize = 30
)

/*
Dispatch basically works like this:
	every m.syncRate tick do
		if bucket needs to sync and isn't syncing
			m.bucketToSyncQueue <- bucket
	each of m.numSyncWorkers do:
		bucket <-m.bucketToSyncQueue
			bucket.sync() with backoff
*/

// A Manager tracks multiple Apigee quotas
type Manager interface {
	Start()
	Apply(auth *auth.Context, p *product.APIProduct, args Args) (*Result, error)
	Close()
}

type manager struct {
	baseURL            *url.URL
	close              chan bool
	client             *http.Client
	now                func() time.Time
	syncRate           time.Duration
	bucketsLock        sync.RWMutex
	buckets            map[string]*bucket // Map from ID -> bucket
	bucketToSyncQueue  chan *bucket
	numSyncWorkers     int
	syncWorkerWG       sync.WaitGroup
	dupCache           ResultCache
	bucketsSyncingLock sync.Mutex
	bucketsSyncing     map[*bucket]struct{}
	org                string
	env                string
	runningContext     context.Context
	cancelContext      context.CancelFunc
}

// NewManager constructs and starts a new Manager. Call Close when done.
func NewManager(options Options) (Manager, error) {
	if err := options.validate(); err != nil {
		return nil, err
	}
	m := newManager(options)
	m.Start()
	return m, nil
}

// newManager constructs a new Manager
func newManager(options Options) *manager {
	return &manager{
		close:             make(chan bool),
		client:            options.Client,
		now:               time.Now,
		syncRate:          defaultSyncRate,
		buckets:           map[string]*bucket{},
		bucketToSyncQueue: make(chan *bucket, syncQueueSize),
		baseURL:           options.BaseURL,
		numSyncWorkers:    defaultNumSyncWorkers,
		dupCache:          ResultCache{size: resultCacheBufferSize},
		bucketsSyncing:    map[*bucket]struct{}{},
		org:               options.Org,
		env:               options.Env,
	}
}

// Start starts the manager.
func (m *manager) Start() {
	log.Infof("starting quota manager")

	m.runningContext, m.cancelContext = context.WithCancel(context.Background())

	go m.bucketMaintenanceLoop()
	for i := 0; i < m.numSyncWorkers; i++ {
		go m.syncBucketDispatcher()
	}

	log.Infof("started quota manager with %d workers", m.numSyncWorkers)
}

// Close shuts down the manager.
func (m *manager) Close() {
	if m == nil {
		return
	}
	log.Infof("closing quota manager")
	m.cancelContext()
	m.runningContext = nil

	m.syncWorkerWG.Add(m.numSyncWorkers)
	m.close <- true
	close(m.bucketToSyncQueue)
	m.syncWorkerWG.Wait()
	log.Infof("closed quota manager")
}

func getQuotaID(auth *auth.Context, p *product.APIProduct) string {
	return fmt.Sprintf("%s-%s", auth.Application, p.Name)
}

// Apply a quota request to the local quota bucket and schedule for sync
func (m *manager) Apply(auth *auth.Context, p *product.APIProduct, args Args) (*Result, error) {

	if result := m.dupCache.Get(args.DeduplicationID); result != nil {
		return result, nil
	}

	quotaID := getQuotaID(auth, p)

	req := &Request{
		Identifier: quotaID,
		Interval:   p.QuotaIntervalInt,
		Allow:      p.QuotaLimitInt,
		TimeUnit:   p.QuotaTimeUnit,
	}

	// a new bucket is created if missing or if product is no longer compatible
	var result *Result
	var err error
	m.bucketsLock.RLock()
	b, ok := m.buckets[quotaID]
	m.bucketsLock.RUnlock()
	if !ok || !b.compatible(req) {
		m.bucketsLock.Lock()
		b, ok = m.buckets[quotaID]
		if !ok || !b.compatible(req) {
			promLabels := m.prometheusLabelsForQuota(quotaID)
			b = newBucket(*req, m, promLabels)
			m.buckets[quotaID] = b
			log.Debugf("new quota bucket: %s", quotaID)
		}
		m.bucketsLock.Unlock()
	}

	req.Weight = args.QuotaAmount
	result, err = b.apply(req)

	if result != nil && err == nil && args.DeduplicationID != "" {
		m.dupCache.Add(args.DeduplicationID, result)
	}

	return result, err
}

// loop to sync active buckets and delete old buckets
func (m *manager) bucketMaintenanceLoop() {
	t := time.NewTicker(m.syncRate)
	for {
		select {
		case <-t.C:
			var deleteIDs []string
			m.bucketsLock.RLock()
			for id, b := range m.buckets {
				if b.needToDelete() {
					deleteIDs = append(deleteIDs, id)

				} else if b.needToSync() {
					m.bucketsSyncingLock.Lock()
					if _, ok := m.bucketsSyncing[b]; !ok { // not already scheduled
						m.bucketsSyncing[b] = struct{}{}
					}
					m.bucketsSyncingLock.Unlock()
					m.bucketToSyncQueue <- b
				}
			}
			m.bucketsLock.RUnlock()

			if deleteIDs != nil {
				log.Debugf("deleting quota buckets: %v", deleteIDs)
				m.bucketsLock.Lock()
				for _, id := range deleteIDs {
					delete(m.buckets, id)
					labels := m.prometheusLabelsForQuota(id)
					prometheusBucketWindowExpires.Delete(labels)
					prometheusBucketChecked.Delete(labels)
					prometheusBucketSynced.Delete(labels)
					prometheusBucketValue.Delete(labels)
				}
				m.bucketsLock.Unlock()
			}

		case <-m.close:
			log.Debugf("closing quota sync loop")
			t.Stop()
			return
		}
	}
}

func (m *manager) prometheusLabelsForQuota(quotaID string) prometheus.Labels {
	return prometheus.Labels{"org": m.org, "env": m.env, "quota": quotaID}
}

// routine for dispatching work to sync a bucket with the server
func (m *manager) syncBucketDispatcher() {

	for bucket := range m.bucketToSyncQueue {
		looper := util.Looper{
			Backoff: util.NewExponentialBackoff(200*time.Millisecond, 30*time.Second, 2, true),
		}
		work := func(ctx context.Context) error {
			return bucket.sync()
		}
		errH := func(err error) error {
			log.Errorf("sync: %s", err)
			return nil
		}

		// run until success or canceled with backoff
		if err := looper.Run(m.runningContext, work, errH); err != nil {
			log.Errorf("looper run: %s", err)
		}

		m.bucketsSyncingLock.Lock()
		delete(m.bucketsSyncing, bucket)
		m.bucketsSyncingLock.Unlock()
	}

	m.syncWorkerWG.Done()
	log.Debugf("closing quota sync worker")
}

// Options allows us to specify options for how this auth manager will run
type Options struct {
	// Client is a configured HTTPClient
	Client *http.Client
	// BaseURL of the Apigee internal proxy
	BaseURL *url.URL
	// Org is organization
	Org string
	// Env is environment
	Env string
}

func (o *Options) validate() error {
	if o.Client == nil ||
		o.BaseURL == nil ||
		o.Org == "" ||
		o.Env == "" {
		return fmt.Errorf("all quota options are required")
	}
	return nil
}

var (
	prometheusBucketValue = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "value",
		Help:      "Current value of a quota",
	}, []string{"org", "env", "quota"})

	prometheusBucketChecked = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "checked",
		Help:      "Time quota was last checked",
	}, []string{"org", "env", "quota"})

	prometheusBucketSynced = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "synced",
		Help:      "Time quota was last synced",
	}, []string{"org", "env", "quota"})

	prometheusBucketWindowExpires = promauto.NewGaugeVec(prometheus.GaugeOpts{
		Subsystem: "quota",
		Name:      "window_expires",
		Help:      "Time quota window will expire",
	}, []string{"org", "env", "quota"})
)
<|endoftext|>


<commit_before>package maas

import (
	"bytes"
	"encoding/base64"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/gomaasapi"
	"launchpad.net/juju-core/environs"
	"math/rand"
	"net/http"
	"net/url"
	"sync"
)

type StorageSuite struct {
	ProviderSuite
}

var _ = Suite(new(StorageSuite))

// makeStorage creates a MAAS storage object for the running test.
func (s *StorageSuite) makeStorage(name string) *maasStorage {
	maasobj := s.testMAASObject.MAASObject
	env := maasEnviron{name: name, maasClientUnlocked: &maasobj}
	return NewStorage(&env).(*maasStorage)
}

// makeRandomBytes returns an array of arbitrary byte values.
func makeRandomBytes(length int) []byte {
	data := make([]byte, length)
	for index := range data {
		data[index] = byte(rand.Intn(256))
	}
	return data
}

// fakeStoredFile creates a file directly in the (simulated) MAAS file store.
// It will contain an arbitrary amount of random data.  The contents are also
// returned.
//
// If you want properly random data here, initialize the randomizer first.
// Or don't, if you want consistent (and debuggable) results.
func (s *StorageSuite) fakeStoredFile(storage environs.Storage, name string) gomaasapi.MAASObject {
	data := makeRandomBytes(rand.Intn(10))
	return s.testMAASObject.TestServer.NewFile(name, data)
}

func (s *StorageSuite) TestGetSnapshotCreatesClone(c *C) {
	original := s.makeStorage("storage-name")
	snapshot := original.getSnapshot()
	c.Check(snapshot.environUnlocked, Equals, original.environUnlocked)
	c.Check(snapshot.maasClientUnlocked.URL().String(), Equals, original.maasClientUnlocked.URL().String())
	// Snapshotting locks the original internally, but does not leave
	// either the original or the snapshot locked.
	c.Check(original.Mutex, Equals, sync.Mutex{})
	c.Check(snapshot.Mutex, Equals, sync.Mutex{})
}

func (s *StorageSuite) TestGetRetrievesFile(c *C) {
	const filename = "stored-data"
	storage := s.makeStorage("get-retrieves-file")
	file := s.fakeStoredFile(storage, filename)
	base64Content, err := file.GetField("content")
	c.Assert(err, IsNil)
	content, err := base64.StdEncoding.DecodeString(base64Content)
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(content))
	c.Check(buf, DeepEquals, content)
}

func (s *StorageSuite) TestRetrieveFileObjectReturnsFileObject(c *C) {
	const filename = "myfile"
	stor := s.makeStorage("rfo-test")
	file := s.fakeStoredFile(stor, filename)
	fileURI, err := file.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	fileContent, err := file.GetField("content")
	c.Assert(err, IsNil)

	obj, err := stor.retrieveFileObject(filename)
	c.Assert(err, IsNil)

	uri, err := obj.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	c.Check(uri, Equals, fileURI)
	content, err := obj.GetField("content")
	c.Check(content, Equals, fileContent)
}

func (s *StorageSuite) TestRetrieveFileObjectReturnsNotFound(c *C) {
	stor := s.makeStorage("rfo-test")
	_, err := stor.retrieveFileObject("nonexistent-file")
	c.Assert(err, NotNil)
	c.Check(err, FitsTypeOf, environs.NotFoundError{})
}

func (s *StorageSuite) TestRetrieveFileObjectEscapesName(c *C) {
	const filename = "#a?b c&d%e!"
	data := []byte("File contents here")
	stor := s.makeStorage("rfo-test")
	err := stor.Put(filename, bytes.NewReader(data), int64(len(data)))
	c.Assert(err, IsNil)

	obj, err := stor.retrieveFileObject(filename)
	c.Assert(err, IsNil)

	base64Content, err := obj.GetField("content")
	c.Assert(err, IsNil)
	content, err := base64.StdEncoding.DecodeString(base64Content)
	c.Assert(err, IsNil)
	c.Check(content, DeepEquals, data)
}

func (s *StorageSuite) TestFileContentsAreBinary(c *C) {
	const filename = "myfile.bin"
	data := []byte{0, 1, 255, 2, 254, 3}
	stor := s.makeStorage("binary-test")

	err := stor.Put(filename, bytes.NewReader(data), int64(len(data)))
	c.Assert(err, IsNil)
	file, err := stor.Get(filename)
	c.Assert(err, IsNil)
	content, err := ioutil.ReadAll(file)
	c.Assert(err, IsNil)

	c.Check(content, DeepEquals, data)
}

func (s *StorageSuite) TestGetReturnsNotFoundErrorIfNotFound(c *C) {
	const filename = "lost-data"
	storage := NewStorage(s.environ)
	_, err := storage.Get(filename)
	c.Assert(err, FitsTypeOf, environs.NotFoundError{})
}

func (s *StorageSuite) TestListReturnsEmptyIfNoFilesStored(c *C) {
	storage := NewStorage(s.environ)
	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestListReturnsAllFilesIfPrefixEmpty(c *C) {
	storage := NewStorage(s.environ)
	files := []string{"1a", "2b", "3c"}
	for _, name := range files {
		s.fakeStoredFile(storage, name)
	}

	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, files)
}

func (s *StorageSuite) TestListSortsResults(c *C) {
	storage := NewStorage(s.environ)
	files := []string{"4d", "1a", "3c", "2b"}
	for _, name := range files {
		s.fakeStoredFile(storage, name)
	}

	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"1a", "2b", "3c", "4d"})
}

func (s *StorageSuite) TestListReturnsNoFilesIfNoFilesMatchPrefix(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "foo")

	listing, err := storage.List("bar")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestListReturnsOnlyFilesWithMatchingPrefix(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "abc")
	s.fakeStoredFile(storage, "xyz")

	listing, err := storage.List("x")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"xyz"})
}

func (s *StorageSuite) TestListMatchesPrefixOnly(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "abc")
	s.fakeStoredFile(storage, "xabc")

	listing, err := storage.List("a")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"abc"})
}

func (s *StorageSuite) TestListOperatesOnFlatNamespace(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "a/b/c/d")

	listing, err := storage.List("a/b")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"a/b/c/d"})
}

// getFileAtURL requests, and returns, the file at the given URL.
func getFileAtURL(fileURL string) ([]byte, error) {
	request, err := http.NewRequest("GET", fileURL, nil)
	if err != nil {
		return nil, err
	}
	client := http.Client{}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

func (s *StorageSuite) TestURLReturnsURLCorrespondingToFile(c *C) {
	const filename = "my-file.txt"
	storage := NewStorage(s.environ).(*maasStorage)
	file := s.fakeStoredFile(storage, filename)
	// The file contains an anon_resource_uri, which lacks a network part
	// (but will probably contain a query part).  anonURL will be the
	// file's full URL.
	anonURI, err := file.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	parsedURI, err := url.Parse(anonURI)
	c.Assert(err, IsNil)
	anonURL := storage.maasClientUnlocked.URL().ResolveReference(parsedURI)
	c.Assert(err, IsNil)

	fileURL, err := storage.URL(filename)
	c.Assert(err, IsNil)

	c.Check(fileURL, NotNil)
	c.Check(fileURL, Equals, anonURL.String())
}

func (s *StorageSuite) TestPutStoresRetrievableFile(c *C) {
	const filename = "broken-toaster.jpg"
	contents := []byte("Contents here")
	length := int64(len(contents))
	storage := NewStorage(s.environ)

	err := storage.Put(filename, bytes.NewReader(contents), length)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(buf, DeepEquals, contents)
}

func (s *StorageSuite) TestPutOverwritesFile(c *C) {
	const filename = "foo.bar"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)
	newContents := []byte("Overwritten")

	err := storage.Put(filename, bytes.NewReader(newContents), int64(len(newContents)))
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(newContents))
	c.Check(buf, DeepEquals, newContents)
}

func (s *StorageSuite) TestPutStopsAtGivenLength(c *C) {
	const filename = "xyzzyz.2.xls"
	const length = 5
	contents := []byte("abcdefghijklmnopqrstuvwxyz")
	storage := NewStorage(s.environ)

	err := storage.Put(filename, bytes.NewReader(contents), length)
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, length)
}

func (s *StorageSuite) TestPutToExistingFileTruncatesAtGivenLength(c *C) {
	const filename = "a-file-which-is-mine"
	oldContents := []byte("abcdefghijklmnopqrstuvwxyz")
	newContents := []byte("xyz")
	storage := NewStorage(s.environ)
	err := storage.Put(filename, bytes.NewReader(oldContents), int64(len(oldContents)))
	c.Assert(err, IsNil)

	err = storage.Put(filename, bytes.NewReader(newContents), int64(len(newContents)))
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(newContents))
	c.Check(buf, DeepEquals, newContents)
}

// These tests not satisfied yet.  Coming in next branch!
/*
func (s *StorageSuite) TestRemoveDeletesFile(c *C) {
	const filename = "doomed.txt"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)

	err := storage.Remove(filename)
	c.Assert(err, IsNil)

	_, err = storage.Get(filename)
	c.Assert(err, FitsTypeOf, environs.NotFoundError{})

	listing, err := storage.List(filename)
	c.Assert(err, IsNil)
	c.Assert(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestRemoveIsIdempotent(c *C) {
	const filename = "half-a-file"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)

	err := storage.Remove(filename)
	c.Assert(err, IsNil)

	err = storage.Remove(filename)
	c.Assert(err, IsNil)
}
*/
<commit_msg>Re-enable tests for file removal in MAAS provider's storage implementation.  Tests: 2 panics.<commit_after>package maas

import (
	"bytes"
	"encoding/base64"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/gomaasapi"
	"launchpad.net/juju-core/environs"
	"math/rand"
	"net/http"
	"net/url"
	"sync"
)

type StorageSuite struct {
	ProviderSuite
}

var _ = Suite(new(StorageSuite))

// makeStorage creates a MAAS storage object for the running test.
func (s *StorageSuite) makeStorage(name string) *maasStorage {
	maasobj := s.testMAASObject.MAASObject
	env := maasEnviron{name: name, maasClientUnlocked: &maasobj}
	return NewStorage(&env).(*maasStorage)
}

// makeRandomBytes returns an array of arbitrary byte values.
func makeRandomBytes(length int) []byte {
	data := make([]byte, length)
	for index := range data {
		data[index] = byte(rand.Intn(256))
	}
	return data
}

// fakeStoredFile creates a file directly in the (simulated) MAAS file store.
// It will contain an arbitrary amount of random data.  The contents are also
// returned.
//
// If you want properly random data here, initialize the randomizer first.
// Or don't, if you want consistent (and debuggable) results.
func (s *StorageSuite) fakeStoredFile(storage environs.Storage, name string) gomaasapi.MAASObject {
	data := makeRandomBytes(rand.Intn(10))
	return s.testMAASObject.TestServer.NewFile(name, data)
}

func (s *StorageSuite) TestGetSnapshotCreatesClone(c *C) {
	original := s.makeStorage("storage-name")
	snapshot := original.getSnapshot()
	c.Check(snapshot.environUnlocked, Equals, original.environUnlocked)
	c.Check(snapshot.maasClientUnlocked.URL().String(), Equals, original.maasClientUnlocked.URL().String())
	// Snapshotting locks the original internally, but does not leave
	// either the original or the snapshot locked.
	c.Check(original.Mutex, Equals, sync.Mutex{})
	c.Check(snapshot.Mutex, Equals, sync.Mutex{})
}

func (s *StorageSuite) TestGetRetrievesFile(c *C) {
	const filename = "stored-data"
	storage := s.makeStorage("get-retrieves-file")
	file := s.fakeStoredFile(storage, filename)
	base64Content, err := file.GetField("content")
	c.Assert(err, IsNil)
	content, err := base64.StdEncoding.DecodeString(base64Content)
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(content))
	c.Check(buf, DeepEquals, content)
}

func (s *StorageSuite) TestRetrieveFileObjectReturnsFileObject(c *C) {
	const filename = "myfile"
	stor := s.makeStorage("rfo-test")
	file := s.fakeStoredFile(stor, filename)
	fileURI, err := file.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	fileContent, err := file.GetField("content")
	c.Assert(err, IsNil)

	obj, err := stor.retrieveFileObject(filename)
	c.Assert(err, IsNil)

	uri, err := obj.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	c.Check(uri, Equals, fileURI)
	content, err := obj.GetField("content")
	c.Check(content, Equals, fileContent)
}

func (s *StorageSuite) TestRetrieveFileObjectReturnsNotFound(c *C) {
	stor := s.makeStorage("rfo-test")
	_, err := stor.retrieveFileObject("nonexistent-file")
	c.Assert(err, NotNil)
	c.Check(err, FitsTypeOf, environs.NotFoundError{})
}

func (s *StorageSuite) TestRetrieveFileObjectEscapesName(c *C) {
	const filename = "#a?b c&d%e!"
	data := []byte("File contents here")
	stor := s.makeStorage("rfo-test")
	err := stor.Put(filename, bytes.NewReader(data), int64(len(data)))
	c.Assert(err, IsNil)

	obj, err := stor.retrieveFileObject(filename)
	c.Assert(err, IsNil)

	base64Content, err := obj.GetField("content")
	c.Assert(err, IsNil)
	content, err := base64.StdEncoding.DecodeString(base64Content)
	c.Assert(err, IsNil)
	c.Check(content, DeepEquals, data)
}

func (s *StorageSuite) TestFileContentsAreBinary(c *C) {
	const filename = "myfile.bin"
	data := []byte{0, 1, 255, 2, 254, 3}
	stor := s.makeStorage("binary-test")

	err := stor.Put(filename, bytes.NewReader(data), int64(len(data)))
	c.Assert(err, IsNil)
	file, err := stor.Get(filename)
	c.Assert(err, IsNil)
	content, err := ioutil.ReadAll(file)
	c.Assert(err, IsNil)

	c.Check(content, DeepEquals, data)
}

func (s *StorageSuite) TestGetReturnsNotFoundErrorIfNotFound(c *C) {
	const filename = "lost-data"
	storage := NewStorage(s.environ)
	_, err := storage.Get(filename)
	c.Assert(err, FitsTypeOf, environs.NotFoundError{})
}

func (s *StorageSuite) TestListReturnsEmptyIfNoFilesStored(c *C) {
	storage := NewStorage(s.environ)
	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestListReturnsAllFilesIfPrefixEmpty(c *C) {
	storage := NewStorage(s.environ)
	files := []string{"1a", "2b", "3c"}
	for _, name := range files {
		s.fakeStoredFile(storage, name)
	}

	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, files)
}

func (s *StorageSuite) TestListSortsResults(c *C) {
	storage := NewStorage(s.environ)
	files := []string{"4d", "1a", "3c", "2b"}
	for _, name := range files {
		s.fakeStoredFile(storage, name)
	}

	listing, err := storage.List("")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"1a", "2b", "3c", "4d"})
}

func (s *StorageSuite) TestListReturnsNoFilesIfNoFilesMatchPrefix(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "foo")

	listing, err := storage.List("bar")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestListReturnsOnlyFilesWithMatchingPrefix(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "abc")
	s.fakeStoredFile(storage, "xyz")

	listing, err := storage.List("x")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"xyz"})
}

func (s *StorageSuite) TestListMatchesPrefixOnly(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "abc")
	s.fakeStoredFile(storage, "xabc")

	listing, err := storage.List("a")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"abc"})
}

func (s *StorageSuite) TestListOperatesOnFlatNamespace(c *C) {
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, "a/b/c/d")

	listing, err := storage.List("a/b")
	c.Assert(err, IsNil)
	c.Check(listing, DeepEquals, []string{"a/b/c/d"})
}

// getFileAtURL requests, and returns, the file at the given URL.
func getFileAtURL(fileURL string) ([]byte, error) {
	request, err := http.NewRequest("GET", fileURL, nil)
	if err != nil {
		return nil, err
	}
	client := http.Client{}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

func (s *StorageSuite) TestURLReturnsURLCorrespondingToFile(c *C) {
	const filename = "my-file.txt"
	storage := NewStorage(s.environ).(*maasStorage)
	file := s.fakeStoredFile(storage, filename)
	// The file contains an anon_resource_uri, which lacks a network part
	// (but will probably contain a query part).  anonURL will be the
	// file's full URL.
	anonURI, err := file.GetField("anon_resource_uri")
	c.Assert(err, IsNil)
	parsedURI, err := url.Parse(anonURI)
	c.Assert(err, IsNil)
	anonURL := storage.maasClientUnlocked.URL().ResolveReference(parsedURI)
	c.Assert(err, IsNil)

	fileURL, err := storage.URL(filename)
	c.Assert(err, IsNil)

	c.Check(fileURL, NotNil)
	c.Check(fileURL, Equals, anonURL.String())
}

func (s *StorageSuite) TestPutStoresRetrievableFile(c *C) {
	const filename = "broken-toaster.jpg"
	contents := []byte("Contents here")
	length := int64(len(contents))
	storage := NewStorage(s.environ)

	err := storage.Put(filename, bytes.NewReader(contents), length)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(buf, DeepEquals, contents)
}

func (s *StorageSuite) TestPutOverwritesFile(c *C) {
	const filename = "foo.bar"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)
	newContents := []byte("Overwritten")

	err := storage.Put(filename, bytes.NewReader(newContents), int64(len(newContents)))
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(newContents))
	c.Check(buf, DeepEquals, newContents)
}

func (s *StorageSuite) TestPutStopsAtGivenLength(c *C) {
	const filename = "xyzzyz.2.xls"
	const length = 5
	contents := []byte("abcdefghijklmnopqrstuvwxyz")
	storage := NewStorage(s.environ)

	err := storage.Put(filename, bytes.NewReader(contents), length)
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, length)
}

func (s *StorageSuite) TestPutToExistingFileTruncatesAtGivenLength(c *C) {
	const filename = "a-file-which-is-mine"
	oldContents := []byte("abcdefghijklmnopqrstuvwxyz")
	newContents := []byte("xyz")
	storage := NewStorage(s.environ)
	err := storage.Put(filename, bytes.NewReader(oldContents), int64(len(oldContents)))
	c.Assert(err, IsNil)

	err = storage.Put(filename, bytes.NewReader(newContents), int64(len(newContents)))
	c.Assert(err, IsNil)

	reader, err := storage.Get(filename)
	c.Assert(err, IsNil)
	defer reader.Close()

	buf, err := ioutil.ReadAll(reader)
	c.Assert(err, IsNil)
	c.Check(len(buf), Equals, len(newContents))
	c.Check(buf, DeepEquals, newContents)
}

func (s *StorageSuite) TestRemoveDeletesFile(c *C) {
	const filename = "doomed.txt"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)

	err := storage.Remove(filename)
	c.Assert(err, IsNil)

	_, err = storage.Get(filename)
	c.Assert(err, FitsTypeOf, environs.NotFoundError{})

	listing, err := storage.List(filename)
	c.Assert(err, IsNil)
	c.Assert(listing, DeepEquals, []string{})
}

func (s *StorageSuite) TestRemoveIsIdempotent(c *C) {
	const filename = "half-a-file"
	storage := NewStorage(s.environ)
	s.fakeStoredFile(storage, filename)

	err := storage.Remove(filename)
	c.Assert(err, IsNil)

	err = storage.Remove(filename)
	c.Assert(err, IsNil)
}
<|endoftext|>


<commit_before>package gossamer

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"strings"
)

func NewMongoStore(hosts string) *MongoStore {
	ms := &MongoStore{
		hosts: hosts,
	}
	ms.Init()

	return ms
}

type MongoStore struct {
	session *mgo.Session
	hosts   string
}

func (m *MongoStore) Init() {
	session, err := mgo.Dial(m.hosts)
	if err != nil {
		log.Println(err)
	}
	session.SetMode(mgo.Monotonic, true)
	m.session = session
}

func (m *MongoStore) Shutdown() {
	m.session.Close()
}

func (m *MongoStore) cloneSession() *mgo.Session {
	return m.session.Clone()
}

func (m *MongoStore) doQuery(query *mgo.Query, ent EntityType, findMultiple bool) interface{} {
	if findMultiple {
		iter := query.Iter()
		switch ent {
		case ENTITY_THINGS:
			rs := []ThingEntity{}
			var r ThingEntity
			for iter.Next(&r) {
				m.postHandleThing(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_OBSERVEDPROPERTIES:
			rs := []ObservedPropertyEntity{}
			var r ObservedPropertyEntity
			for iter.Next(&r) {
				m.postHandleObservedProperty(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_LOCATIONS:
			rs := []LocationEntity{}
			var r LocationEntity
			for iter.Next(&r) {
				m.postHandleLocation(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_DATASTREAMS:
			rs := []DatastreamEntity{}
			var r DatastreamEntity
			for iter.Next(&r) {
				m.postHandleDatastream(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_SENSORS:
			rs := []SensorEntity{}
			var r SensorEntity
			for iter.Next(&r) {
				m.postHandleSensor(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_OBSERVATIONS:
			rs := []ObservationEntity{}
			var r ObservationEntity
			for iter.Next(&r) {
				m.postHandleObservation(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_FEATURESOFINTERESTS:
			rs := []FeatureOfInterestEntity{}
			var r FeatureOfInterestEntity
			for iter.Next(&r) {
				m.postHandleFeatureOfInterest(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_HISTORICALLOCATIONS:
			rs := []HistoricalLocationEntity{}
			var r HistoricalLocationEntity
			for iter.Next(&r) {
				m.postHandleHistoricalLocation(&r)
				rs = append(rs, r)
			}
			return rs
		}
	} else {
		switch {
		case ent == ENTITY_THINGS || ent == ENTITY_THING:
			var r ThingEntity
			query.One(&r)
			m.postHandleThing(&r)
			return r

		case ent == ENTITY_OBSERVEDPROPERTIES || ent == ENTITY_OBSERVEDPROPERTY:
			var r ObservedPropertyEntity
			query.One(&r)
			m.postHandleObservedProperty(&r)
			return r

		case ent == ENTITY_LOCATIONS || ent == ENTITY_LOCATION:
			var r LocationEntity
			query.One(&r)
			m.postHandleLocation(&r)
			return r

		case ent == ENTITY_DATASTREAMS || ent == ENTITY_DATASTREAM:
			var r DatastreamEntity
			query.One(&r)
			m.postHandleDatastream(&r)
			return r

		case ent == ENTITY_SENSORS || ent == ENTITY_SENSOR:
			var r SensorEntity
			query.One(&r)
			m.postHandleSensor(&r)
			return r

		case ent == ENTITY_OBSERVATIONS || ent == ENTITY_OBSERVATION:
			var r ObservationEntity
			query.One(&r)
			m.postHandleObservation(&r)
			return r

		case ent == ENTITY_FEATURESOFINTERESTS || ent == ENTITY_FEATURESOFINTEREST:
			var r FeatureOfInterestEntity
			query.One(&r)
			m.postHandleFeatureOfInterest(&r)
			return r

		case ent == ENTITY_HISTORICALLOCATIONS || ent == ENTITY_HISTORICALLOCATION:
			var r HistoricalLocationEntity
			query.One(&r)
			m.postHandleHistoricalLocation(&r)
			return r
		}
	}
	return nil
}

func (m *MongoStore) createQuery(c *mgo.Collection, rp ResourcePath, lastResult interface{}) (query *mgo.Query, findMultiple bool) {
	last := rp.At(rp.CurrentIndex() - 1)
	curr := rp.At(rp.CurrentIndex())
	currEntity := curr.GetEntity()
	isFirst := rp.IsFirst()

	bsonMap := bson.M{}
	if curr.GetId() != "" {
		bsonMap["@iot_id"] = curr.GetId()
		findMultiple = false
	} else if IsSingularEntity(string(currEntity)) {
		log.Println("Is Singular")
		bsonMap["@iot_id"] = lastResult.(SensorThing).GetAssociatedEntityId(currEntity)
		findMultiple = false
	} else {
		findMultiple = true
		if !isFirst {
			lastEntity := last.GetEntity()

			switch {
			case lastEntity == ENTITY_LOCATIONS && currEntity == ENTITY_THINGS:
				// Locations(id)/Things
				// MULTIPLE: { where iot_locations_id contains prev.id}
				break

			case lastEntity == ENTITY_LOCATIONS && currEntity == ENTITY_HISTORICALLOCATIONS:
				//				# Locations(id)/HistoricalLocations
				//				MULTIPLE: { where iot_locations_id contains prev.id}
				break

			case lastEntity == ENTITY_HISTORICALLOCATIONS && currEntity == ENTITY_LOCATIONS:
				// # HistoricalLocations(id)/Locations
				// MULTIPLE: { where iot_id contains prev.id }
				break

			case lastEntity == ENTITY_THINGS && currEntity == ENTITY_LOCATIONS:
				// # Things(id)/Locations
				// MULTIPLE: { where iot_id is in prev.iot_locations_id}
				break

			default:
				bsonMap["@iot_"+strings.ToLower(string(lastEntity))+"_id"] = last.GetId()
				break
			}
		} else {
			log.Println("Uh oh..")
		}
	}
	log.Println("bsonMap", bsonMap)
	query = c.Find(bsonMap)

	return
}

func (m *MongoStore) Query(rp ResourcePath) (interface{}, error) {
	queryComplete := make(chan bool)
	var results interface{}
	go func() {
		session := m.cloneSession()
		defer session.Close()

		for rp.HasNext() {
			curr := rp.Next()
			currEntity := curr.GetEntity()

			c := session.DB("sensorthings").C(ResolveMongoCollectionName(currEntity))
			query, findMultiple := m.createQuery(c, rp, results)

			resourceQueryComplete := make(chan bool)
			go func() {
				results = m.doQuery(query, currEntity, findMultiple)
				resourceQueryComplete <- true
			}()
			<-resourceQueryComplete

			if rp.IsLast() {
				break
			}
		}
		queryComplete <- true
	}()
	<-queryComplete
	return results, nil
}

func (m *MongoStore) postHandleThing(e *ThingEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_THINGS)

	e.NavLinkLocations = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/Locations"
	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/HistoricalLocations"
	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/Datastreams"
}

func (m *MongoStore) postHandleObservedProperty(e *ObservedPropertyEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_OBSERVEDPROPERTIES)

	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_OBSERVEDPROPERTIES) + "/Datastreams"
}

func (m *MongoStore) postHandleLocation(e *LocationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_LOCATIONS)

	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_LOCATIONS) + "/HistoricalLocations"
	e.NavLinkThings = ResolveEntityLink(e.Id, ENTITY_LOCATIONS) + "/Things"
}

func (m *MongoStore) postHandleDatastream(e *DatastreamEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_DATASTREAMS)

	e.NavLinkObservations = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Observations"
	e.NavLinkObservedProperty = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/ObservedProperty"
	e.NavLinkSensor = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Sensor"
	e.NavLinkThing = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Thing"
}

func (m *MongoStore) postHandleSensor(e *SensorEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_SENSORS)

	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_SENSORS) + "/Datastreams"
}

func (m *MongoStore) postHandleObservation(e *ObservationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_OBSERVATIONS)

	e.NavLinkDatastream = ResolveEntityLink(e.Id, ENTITY_OBSERVATIONS) + "/Datastream"
	e.NavLinkFeatureOfInterest = ResolveEntityLink(e.Id, ENTITY_OBSERVATIONS) + "/FeatureOfInterest"
}

func (m *MongoStore) postHandleFeatureOfInterest(e *FeatureOfInterestEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_FEATURESOFINTERESTS)

	e.NavLinkObservations = ResolveEntityLink(e.Id, ENTITY_FEATURESOFINTERESTS) + "/Observations"
}

func (m *MongoStore) postHandleHistoricalLocation(e *HistoricalLocationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_HISTORICALLOCATIONS)

	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_HISTORICALLOCATIONS) + "/HistoricalLocations"
	e.NavLinkThing = ResolveEntityLink(e.Id, ENTITY_HISTORICALLOCATIONS) + "/Thing"
}

func ResolveMongoCollectionName(ent EntityType) string {
	switch {
	case ent == ENTITY_THINGS || ent == ENTITY_THING:
		return "things"

	case ent == ENTITY_OBSERVEDPROPERTY || ent == ENTITY_OBSERVEDPROPERTIES:
		return "observedproperties"

	case ent == ENTITY_LOCATION || ent == ENTITY_LOCATIONS:
		return "locations"

	case ent == ENTITY_DATASTREAM || ent == ENTITY_DATASTREAMS:
		return "datastreams"

	case ent == ENTITY_SENSOR || ent == ENTITY_SENSORS:
		return "sensors"

	case ent == ENTITY_OBSERVATION || ent == ENTITY_OBSERVATIONS:
		return "observations"

	case ent == ENTITY_FEATURESOFINTEREST || ent == ENTITY_FEATURESOFINTERESTS:
		return "featureofinterests"

	case ent == ENTITY_HISTORICALLOCATION || ent == ENTITY_HISTORICALLOCATIONS:
		return "historicallocations"
	}
	return ""
}
<commit_msg>Many-Many relation between Locations and Things, Locations and HistoricalLocations etc (Mongo Store)<commit_after>package gossamer

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"strings"
)

func NewMongoStore(hosts string) *MongoStore {
	ms := &MongoStore{
		hosts: hosts,
	}
	ms.Init()

	return ms
}

type MongoStore struct {
	session *mgo.Session
	hosts   string
}

func (m *MongoStore) Init() {
	session, err := mgo.Dial(m.hosts)
	if err != nil {
		log.Println(err)
	}
	session.SetMode(mgo.Monotonic, true)
	m.session = session
}

func (m *MongoStore) Shutdown() {
	m.session.Close()
}

func (m *MongoStore) cloneSession() *mgo.Session {
	return m.session.Clone()
}

func (m *MongoStore) doQuery(query *mgo.Query, ent EntityType, findMultiple bool) interface{} {
	if findMultiple {
		iter := query.Iter()
		switch ent {
		case ENTITY_THINGS:
			rs := []ThingEntity{}
			var r ThingEntity
			for iter.Next(&r) {
				m.postHandleThing(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_OBSERVEDPROPERTIES:
			rs := []ObservedPropertyEntity{}
			var r ObservedPropertyEntity
			for iter.Next(&r) {
				m.postHandleObservedProperty(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_LOCATIONS:
			rs := []LocationEntity{}
			var r LocationEntity
			for iter.Next(&r) {
				m.postHandleLocation(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_DATASTREAMS:
			rs := []DatastreamEntity{}
			var r DatastreamEntity
			for iter.Next(&r) {
				m.postHandleDatastream(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_SENSORS:
			rs := []SensorEntity{}
			var r SensorEntity
			for iter.Next(&r) {
				m.postHandleSensor(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_OBSERVATIONS:
			rs := []ObservationEntity{}
			var r ObservationEntity
			for iter.Next(&r) {
				m.postHandleObservation(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_FEATURESOFINTERESTS:
			rs := []FeatureOfInterestEntity{}
			var r FeatureOfInterestEntity
			for iter.Next(&r) {
				m.postHandleFeatureOfInterest(&r)
				rs = append(rs, r)
			}
			return rs

		case ENTITY_HISTORICALLOCATIONS:
			rs := []HistoricalLocationEntity{}
			var r HistoricalLocationEntity
			for iter.Next(&r) {
				m.postHandleHistoricalLocation(&r)
				rs = append(rs, r)
			}
			return rs
		}
	} else {
		switch {
		case ent == ENTITY_THINGS || ent == ENTITY_THING:
			var r ThingEntity
			query.One(&r)
			m.postHandleThing(&r)
			return r

		case ent == ENTITY_OBSERVEDPROPERTIES || ent == ENTITY_OBSERVEDPROPERTY:
			var r ObservedPropertyEntity
			query.One(&r)
			m.postHandleObservedProperty(&r)
			return r

		case ent == ENTITY_LOCATIONS || ent == ENTITY_LOCATION:
			var r LocationEntity
			query.One(&r)
			m.postHandleLocation(&r)
			return r

		case ent == ENTITY_DATASTREAMS || ent == ENTITY_DATASTREAM:
			var r DatastreamEntity
			query.One(&r)
			m.postHandleDatastream(&r)
			return r

		case ent == ENTITY_SENSORS || ent == ENTITY_SENSOR:
			var r SensorEntity
			query.One(&r)
			m.postHandleSensor(&r)
			return r

		case ent == ENTITY_OBSERVATIONS || ent == ENTITY_OBSERVATION:
			var r ObservationEntity
			query.One(&r)
			m.postHandleObservation(&r)
			return r

		case ent == ENTITY_FEATURESOFINTERESTS || ent == ENTITY_FEATURESOFINTEREST:
			var r FeatureOfInterestEntity
			query.One(&r)
			m.postHandleFeatureOfInterest(&r)
			return r

		case ent == ENTITY_HISTORICALLOCATIONS || ent == ENTITY_HISTORICALLOCATION:
			var r HistoricalLocationEntity
			query.One(&r)
			m.postHandleHistoricalLocation(&r)
			return r
		}
	}
	return nil
}

func (m *MongoStore) createQuery(c *mgo.Collection, rp ResourcePath, lastResult interface{}) (query *mgo.Query, findMultiple bool) {
	last := rp.At(rp.CurrentIndex() - 1)
	curr := rp.At(rp.CurrentIndex())
	currEntity := curr.GetEntity()
	isFirst := rp.IsFirst()

	bsonMap := bson.M{}
	if curr.GetId() != "" {
		bsonMap["@iot_id"] = curr.GetId()
		findMultiple = false
	} else if IsSingularEntity(string(currEntity)) {
		bsonMap["@iot_id"] = lastResult.(SensorThing).GetAssociatedEntityId(currEntity)
		findMultiple = false
	} else {
		findMultiple = true
		if !isFirst {
			lastEntity := last.GetEntity()

			switch {
			case lastEntity == ENTITY_LOCATIONS && currEntity == ENTITY_THINGS,
				 lastEntity == ENTITY_LOCATIONS && currEntity == ENTITY_HISTORICALLOCATIONS:

				bsonMap["@iot_locations_id"] = map[string]interface{}{
					"$in": []string{last.GetId()},
				}

			case lastEntity == ENTITY_HISTORICALLOCATIONS && currEntity == ENTITY_LOCATIONS:
				bsonMap["@iot_id"] = map[string]interface{}{
					"$in": lastResult.(HistoricalLocationEntity).IdLocations,
				}

			case lastEntity == ENTITY_THINGS && currEntity == ENTITY_LOCATIONS:
				log.Println("Case D")
				bsonMap["@iot_id"] = map[string]interface{}{
					"$in": lastResult.(ThingEntity).IdLocations,
				}

			default:
				bsonMap["@iot_"+strings.ToLower(string(lastEntity))+"_id"] = last.GetId()
			}
		}
	}
	query = c.Find(bsonMap)

	return
}

func (m *MongoStore) Query(rp ResourcePath) (interface{}, error) {
	queryComplete := make(chan bool)
	var results interface{}
	go func() {
		session := m.cloneSession()
		defer session.Close()

		for rp.HasNext() {
			curr := rp.Next()
			currEntity := curr.GetEntity()

			c := session.DB("sensorthings").C(ResolveMongoCollectionName(currEntity))
			query, findMultiple := m.createQuery(c, rp, results)

			resourceQueryComplete := make(chan bool)
			go func() {
				results = m.doQuery(query, currEntity, findMultiple)
				resourceQueryComplete <- true
			}()
			<-resourceQueryComplete

			if rp.IsLast() {
				break
			}
		}
		queryComplete <- true
	}()
	<-queryComplete
	return results, nil
}

func (m *MongoStore) postHandleThing(e *ThingEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_THINGS)

	e.NavLinkLocations = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/Locations"
	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/HistoricalLocations"
	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_THINGS) + "/Datastreams"
}

func (m *MongoStore) postHandleObservedProperty(e *ObservedPropertyEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_OBSERVEDPROPERTIES)

	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_OBSERVEDPROPERTIES) + "/Datastreams"
}

func (m *MongoStore) postHandleLocation(e *LocationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_LOCATIONS)

	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_LOCATIONS) + "/HistoricalLocations"
	e.NavLinkThings = ResolveEntityLink(e.Id, ENTITY_LOCATIONS) + "/Things"
}

func (m *MongoStore) postHandleDatastream(e *DatastreamEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_DATASTREAMS)

	e.NavLinkObservations = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Observations"
	e.NavLinkObservedProperty = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/ObservedProperty"
	e.NavLinkSensor = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Sensor"
	e.NavLinkThing = ResolveEntityLink(e.Id, ENTITY_DATASTREAMS) + "/Thing"
}

func (m *MongoStore) postHandleSensor(e *SensorEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_SENSORS)

	e.NavLinkDatastreams = ResolveEntityLink(e.Id, ENTITY_SENSORS) + "/Datastreams"
}

func (m *MongoStore) postHandleObservation(e *ObservationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_OBSERVATIONS)

	e.NavLinkDatastream = ResolveEntityLink(e.Id, ENTITY_OBSERVATIONS) + "/Datastream"
	e.NavLinkFeatureOfInterest = ResolveEntityLink(e.Id, ENTITY_OBSERVATIONS) + "/FeatureOfInterest"
}

func (m *MongoStore) postHandleFeatureOfInterest(e *FeatureOfInterestEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_FEATURESOFINTERESTS)

	e.NavLinkObservations = ResolveEntityLink(e.Id, ENTITY_FEATURESOFINTERESTS) + "/Observations"
}

func (m *MongoStore) postHandleHistoricalLocation(e *HistoricalLocationEntity) {
	e.SelfLink = ResolveSelfLinkUrl(e.Id, ENTITY_HISTORICALLOCATIONS)

	e.NavLinkHistoricalLocations = ResolveEntityLink(e.Id, ENTITY_HISTORICALLOCATIONS) + "/HistoricalLocations"
	e.NavLinkThing = ResolveEntityLink(e.Id, ENTITY_HISTORICALLOCATIONS) + "/Thing"
}

func ResolveMongoCollectionName(ent EntityType) string {
	switch {
	case ent == ENTITY_THINGS || ent == ENTITY_THING:
		return "things"

	case ent == ENTITY_OBSERVEDPROPERTY || ent == ENTITY_OBSERVEDPROPERTIES:
		return "observedproperties"

	case ent == ENTITY_LOCATION || ent == ENTITY_LOCATIONS:
		return "locations"

	case ent == ENTITY_DATASTREAM || ent == ENTITY_DATASTREAMS:
		return "datastreams"

	case ent == ENTITY_SENSOR || ent == ENTITY_SENSORS:
		return "sensors"

	case ent == ENTITY_OBSERVATION || ent == ENTITY_OBSERVATIONS:
		return "observations"

	case ent == ENTITY_FEATURESOFINTEREST || ent == ENTITY_FEATURESOFINTERESTS:
		return "featureofinterests"

	case ent == ENTITY_HISTORICALLOCATION || ent == ENTITY_HISTORICALLOCATIONS:
		return "historicallocations"
	}
	return ""
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package awstasks

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"

	"k8s.io/apimachinery/pkg/util/validation/field"
	"k8s.io/klog/v2"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraformWriter"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/fi/cloudup/cloudformation"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
)

// +kops:fitask
type SQS struct {
	Name      *string
	Lifecycle fi.Lifecycle

	ARN                    *string
	URL                    *string
	MessageRetentionPeriod int
	Policy                 fi.Resource

	Tags map[string]string
}

var _ fi.CompareWithID = &SQS{}

func (q *SQS) CompareWithID() *string {
	return q.ARN
}

func (q *SQS) Find(c *fi.Context) (*SQS, error) {
	cloud := c.Cloud.(awsup.AWSCloud)

	if q.Name == nil {
		return nil, nil
	}

	response, err := cloud.SQS().ListQueues(&sqs.ListQueuesInput{
		MaxResults:      aws.Int64(2),
		QueueNamePrefix: q.Name,
	})
	if err != nil {
		return nil, fmt.Errorf("error listing SQS queues: %v", err)
	}
	if response == nil || len(response.QueueUrls) == 0 {
		return nil, nil
	}
	if len(response.QueueUrls) != 1 {
		return nil, fmt.Errorf("found multiple SQS queues matching queue name")
	}
	url := response.QueueUrls[0]

	attributes, err := cloud.SQS().GetQueueAttributes(&sqs.GetQueueAttributesInput{
		AttributeNames: []*string{s("MessageRetentionPeriod"), s("Policy"), s("QueueArn")},
		QueueUrl:       url,
	})
	if err != nil {
		return nil, fmt.Errorf("error getting SQS queue attributes: %v", err)
	}
	actualPolicy := *attributes.Attributes["Policy"]
	actualARN := *attributes.Attributes["QueueArn"]
	period, err := strconv.Atoi(*attributes.Attributes["MessageRetentionPeriod"])
	if err != nil {
		return nil, fmt.Errorf("error coverting MessageRetentionPeriod to int: %v", err)
	}

	tags, err := cloud.SQS().ListQueueTags(&sqs.ListQueueTagsInput{
		QueueUrl: url,
	})
	if err != nil {
		return nil, fmt.Errorf("error listing SQS queue tags: %v", err)
	}

	// We parse both as JSON; if the json forms are equal we pretend the actual value is the expected value
	if q.Policy != nil {
		expectedPolicy, err := fi.ResourceAsString(q.Policy)
		if err != nil {
			return nil, fmt.Errorf("error reading expected Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}
		expectedJson := make(map[string]interface{})
		err = json.Unmarshal([]byte(expectedPolicy), &expectedJson)
		if err != nil {
			return nil, fmt.Errorf("error parsing expected Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}
		actualJson := make(map[string]interface{})
		err = json.Unmarshal([]byte(actualPolicy), &actualJson)
		if err != nil {
			return nil, fmt.Errorf("error parsing actual Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}

		if reflect.DeepEqual(actualJson, expectedJson) {
			klog.V(2).Infof("actual Policy was json-equal to expected; returning expected value")
			actualPolicy = expectedPolicy
		}
	}

	actual := &SQS{
		ARN:                    s(actualARN),
		Name:                   q.Name,
		URL:                    url,
		Lifecycle:              q.Lifecycle,
		Policy:                 fi.NewStringResource(actualPolicy),
		MessageRetentionPeriod: period,
		Tags:                   intersectSQSTags(tags.Tags, q.Tags),
	}

	// Avoid flapping
	q.ARN = actual.ARN

	return actual, nil
}

func (q *SQS) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(q, c)
}

func (q *SQS) CheckChanges(a, e, changes *SQS) error {
	if a == nil {
		if e.Name == nil {
			return field.Required(field.NewPath("Name"), "")
		}
	}
	if a != nil {
		if changes.URL != nil {
			return fi.CannotChangeField("URL")
		}
	}
	return nil
}

func (q *SQS) RenderAWS(t *awsup.AWSAPITarget, a, e, changes *SQS) error {
	policy, err := fi.ResourceAsString(e.Policy)
	if err != nil {
		return fmt.Errorf("error rendering RolePolicyDocument: %v", err)
	}

	if a == nil {
		request := &sqs.CreateQueueInput{
			Attributes: map[string]*string{
				"MessageRetentionPeriod": s(strconv.Itoa(q.MessageRetentionPeriod)),
				"Policy":                 s(policy),
			},
			QueueName: q.Name,
			Tags:      convertTagsToPointers(q.Tags),
		}
		response, err := t.Cloud.SQS().CreateQueue(request)
		if err != nil {
			return fmt.Errorf("error creating SQS queue: %v", err)
		}

		attributes, err := t.Cloud.SQS().GetQueueAttributes(&sqs.GetQueueAttributesInput{
			AttributeNames: []*string{s("QueueArn")},
			QueueUrl:       response.QueueUrl,
		})
		if err != nil {
			return fmt.Errorf("error getting SQS queue attributes: %v", err)
		}

		e.ARN = attributes.Attributes["QueueArn"]
	}

	return nil
}

type terraformSQSQueue struct {
	Name                    *string                  `cty:"name"`
	MessageRetentionSeconds int                      `cty:"message_retention_seconds"`
	Policy                  *terraformWriter.Literal `cty:"policy"`
	Tags                    map[string]string        `cty:"tags"`
}

func (_ *SQS) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *SQS) error {
	p, err := t.AddFileResource("aws_sqs_queue", *e.Name, "policy", e.Policy, false)
	if err != nil {
		return err
	}

	tf := &terraformSQSQueue{
		Name:                    e.Name,
		MessageRetentionSeconds: e.MessageRetentionPeriod,
		Policy:                  p,
		Tags:                    e.Tags,
	}

	return t.RenderResource("aws_sqs_queue", *e.Name, tf)
}

func (e *SQS) TerraformLink() *terraformWriter.Literal {
	return terraformWriter.LiteralProperty("aws_sqs_queue", *e.Name, "arn")
}

type cloudformationSQSQueue struct {
	QueueName              *string             `json:"QueueName"`
	MessageRetentionPeriod int                 `json:"MessageRetentionPeriod"`
	Tags                   []cloudformationTag `json:"Tags,omitempty"`
}

type cloudformationSQSQueuePolicy struct {
	Queues         []*cloudformation.Literal `json:"Queues"`
	PolicyDocument map[string]interface{}    `json:"PolicyDocument"`
	Tags           []cloudformationTag       `json:"Tags,omitempty"`
}

func (_ *SQS) RenderCloudformation(t *cloudformation.CloudformationTarget, a, e, changes *SQS) error {
	cfQueue := &cloudformationSQSQueue{
		QueueName:              e.Name,
		MessageRetentionPeriod: e.MessageRetentionPeriod,
		Tags:                   buildCloudformationTags(e.Tags),
	}

	err := t.RenderResource("AWS::SQS::Queue", *e.Name, cfQueue)
	if err != nil {
		return err
	}

	// convert Policy string into json
	jsonString, err := fi.ResourceAsBytes(e.Policy)
	if err != nil {
		return err
	}
	data := make(map[string]interface{})
	err = json.Unmarshal(jsonString, &data)
	if err != nil {
		return fmt.Errorf("error parsing SQS PolicyDocument: %v", err)
	}

	cfQueueRef := cloudformation.Ref("AWS::SQS::Queue", fi.StringValue(e.Name))

	cfQueuePolicy := &cloudformationSQSQueuePolicy{
		Queues:         []*cloudformation.Literal{cfQueueRef},
		PolicyDocument: data,
	}
	return t.RenderResource("AWS::SQS::QueuePolicy", *e.Name+"Policy", cfQueuePolicy)
}

func (e *SQS) CloudformationLink() *cloudformation.Literal {
	return cloudformation.Ref("AWS::SQS::Queue", *e.Name)
}

// change tags to format required by CreateQueue
func convertTagsToPointers(tags map[string]string) map[string]*string {
	newTags := map[string]*string{}
	for k, v := range tags {
		vv := v
		newTags[k] = &vv
	}

	return newTags
}

// intersectSQSTags does the same thing as intersectTags, but takes different input because SQS tags are listed differently
func intersectSQSTags(tags map[string]*string, desired map[string]string) map[string]string {
	if tags == nil {
		return nil
	}
	actual := make(map[string]string)
	for k, v := range tags {
		vv := aws.StringValue(v)

		if _, found := desired[k]; found {
			actual[k] = vv
		}
	}
	if len(actual) == 0 && desired == nil {
		// Avoid problems with comparison between nil & {}
		return nil
	}
	return actual
}
<commit_msg>Consistent formating of SQS policy to prevent flapping<commit_after>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package awstasks

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"

	"k8s.io/apimachinery/pkg/util/validation/field"
	"k8s.io/klog/v2"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraformWriter"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/fi/cloudup/cloudformation"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
)

// +kops:fitask
type SQS struct {
	Name      *string
	Lifecycle fi.Lifecycle

	ARN                    *string
	URL                    *string
	MessageRetentionPeriod int
	Policy                 fi.Resource

	Tags map[string]string
}

var _ fi.CompareWithID = &SQS{}

func (q *SQS) CompareWithID() *string {
	return q.ARN
}

func (q *SQS) Find(c *fi.Context) (*SQS, error) {
	cloud := c.Cloud.(awsup.AWSCloud)

	if q.Name == nil {
		return nil, nil
	}

	response, err := cloud.SQS().ListQueues(&sqs.ListQueuesInput{
		MaxResults:      aws.Int64(2),
		QueueNamePrefix: q.Name,
	})
	if err != nil {
		return nil, fmt.Errorf("error listing SQS queues: %v", err)
	}
	if response == nil || len(response.QueueUrls) == 0 {
		return nil, nil
	}
	if len(response.QueueUrls) != 1 {
		return nil, fmt.Errorf("found multiple SQS queues matching queue name")
	}
	url := response.QueueUrls[0]

	attributes, err := cloud.SQS().GetQueueAttributes(&sqs.GetQueueAttributesInput{
		AttributeNames: []*string{s("MessageRetentionPeriod"), s("Policy"), s("QueueArn")},
		QueueUrl:       url,
	})
	if err != nil {
		return nil, fmt.Errorf("error getting SQS queue attributes: %v", err)
	}
	actualPolicy := *attributes.Attributes["Policy"]
	actualARN := *attributes.Attributes["QueueArn"]
	period, err := strconv.Atoi(*attributes.Attributes["MessageRetentionPeriod"])
	if err != nil {
		return nil, fmt.Errorf("error coverting MessageRetentionPeriod to int: %v", err)
	}

	tags, err := cloud.SQS().ListQueueTags(&sqs.ListQueueTagsInput{
		QueueUrl: url,
	})
	if err != nil {
		return nil, fmt.Errorf("error listing SQS queue tags: %v", err)
	}

	// We parse both as JSON; if the json forms are equal we pretend the actual value is the expected value
	if q.Policy != nil {
		expectedPolicy, err := fi.ResourceAsString(q.Policy)
		if err != nil {
			return nil, fmt.Errorf("error reading expected Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}
		expectedJson := make(map[string]interface{})
		err = json.Unmarshal([]byte(expectedPolicy), &expectedJson)
		if err != nil {
			return nil, fmt.Errorf("error parsing expected Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}
		actualJson := make(map[string]interface{})
		err = json.Unmarshal([]byte(actualPolicy), &actualJson)
		if err != nil {
			return nil, fmt.Errorf("error parsing actual Policy for SQS %q: %v", aws.StringValue(q.Name), err)
		}

		if reflect.DeepEqual(actualJson, expectedJson) {
			klog.V(2).Infof("actual Policy was json-equal to expected; returning expected value")
			actualPolicy = expectedPolicy
			q.Policy = fi.NewStringResource(expectedPolicy)
		}

	}

	actual := &SQS{
		ARN:                    s(actualARN),
		Name:                   q.Name,
		URL:                    url,
		Lifecycle:              q.Lifecycle,
		Policy:                 fi.NewStringResource(actualPolicy),
		MessageRetentionPeriod: period,
		Tags:                   intersectSQSTags(tags.Tags, q.Tags),
	}

	// Avoid flapping
	q.ARN = actual.ARN

	return actual, nil
}

func (q *SQS) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(q, c)
}

func (q *SQS) CheckChanges(a, e, changes *SQS) error {
	if a == nil {
		if e.Name == nil {
			return field.Required(field.NewPath("Name"), "")
		}
	}
	if a != nil {
		if changes.URL != nil {
			return fi.CannotChangeField("URL")
		}
	}
	return nil
}

func (q *SQS) RenderAWS(t *awsup.AWSAPITarget, a, e, changes *SQS) error {
	policy, err := fi.ResourceAsString(e.Policy)
	if err != nil {
		return fmt.Errorf("error rendering RolePolicyDocument: %v", err)
	}

	if a == nil {
		request := &sqs.CreateQueueInput{
			Attributes: map[string]*string{
				"MessageRetentionPeriod": s(strconv.Itoa(q.MessageRetentionPeriod)),
				"Policy":                 s(policy),
			},
			QueueName: q.Name,
			Tags:      convertTagsToPointers(q.Tags),
		}
		response, err := t.Cloud.SQS().CreateQueue(request)
		if err != nil {
			return fmt.Errorf("error creating SQS queue: %v", err)
		}

		attributes, err := t.Cloud.SQS().GetQueueAttributes(&sqs.GetQueueAttributesInput{
			AttributeNames: []*string{s("QueueArn")},
			QueueUrl:       response.QueueUrl,
		})
		if err != nil {
			return fmt.Errorf("error getting SQS queue attributes: %v", err)
		}

		e.ARN = attributes.Attributes["QueueArn"]
	}

	return nil
}

type terraformSQSQueue struct {
	Name                    *string                  `cty:"name"`
	MessageRetentionSeconds int                      `cty:"message_retention_seconds"`
	Policy                  *terraformWriter.Literal `cty:"policy"`
	Tags                    map[string]string        `cty:"tags"`
}

func (_ *SQS) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *SQS) error {
	p, err := t.AddFileResource("aws_sqs_queue", *e.Name, "policy", e.Policy, false)
	if err != nil {
		return err
	}

	tf := &terraformSQSQueue{
		Name:                    e.Name,
		MessageRetentionSeconds: e.MessageRetentionPeriod,
		Policy:                  p,
		Tags:                    e.Tags,
	}

	return t.RenderResource("aws_sqs_queue", *e.Name, tf)
}

func (e *SQS) TerraformLink() *terraformWriter.Literal {
	return terraformWriter.LiteralProperty("aws_sqs_queue", *e.Name, "arn")
}

type cloudformationSQSQueue struct {
	QueueName              *string             `json:"QueueName"`
	MessageRetentionPeriod int                 `json:"MessageRetentionPeriod"`
	Tags                   []cloudformationTag `json:"Tags,omitempty"`
}

type cloudformationSQSQueuePolicy struct {
	Queues         []*cloudformation.Literal `json:"Queues"`
	PolicyDocument map[string]interface{}    `json:"PolicyDocument"`
	Tags           []cloudformationTag       `json:"Tags,omitempty"`
}

func (_ *SQS) RenderCloudformation(t *cloudformation.CloudformationTarget, a, e, changes *SQS) error {
	cfQueue := &cloudformationSQSQueue{
		QueueName:              e.Name,
		MessageRetentionPeriod: e.MessageRetentionPeriod,
		Tags:                   buildCloudformationTags(e.Tags),
	}

	err := t.RenderResource("AWS::SQS::Queue", *e.Name, cfQueue)
	if err != nil {
		return err
	}

	// convert Policy string into json
	jsonString, err := fi.ResourceAsBytes(e.Policy)
	if err != nil {
		return err
	}
	data := make(map[string]interface{})
	err = json.Unmarshal(jsonString, &data)
	if err != nil {
		return fmt.Errorf("error parsing SQS PolicyDocument: %v", err)
	}

	cfQueueRef := cloudformation.Ref("AWS::SQS::Queue", fi.StringValue(e.Name))

	cfQueuePolicy := &cloudformationSQSQueuePolicy{
		Queues:         []*cloudformation.Literal{cfQueueRef},
		PolicyDocument: data,
	}
	return t.RenderResource("AWS::SQS::QueuePolicy", *e.Name+"Policy", cfQueuePolicy)
}

func (e *SQS) CloudformationLink() *cloudformation.Literal {
	return cloudformation.Ref("AWS::SQS::Queue", *e.Name)
}

// change tags to format required by CreateQueue
func convertTagsToPointers(tags map[string]string) map[string]*string {
	newTags := map[string]*string{}
	for k, v := range tags {
		vv := v
		newTags[k] = &vv
	}

	return newTags
}

// intersectSQSTags does the same thing as intersectTags, but takes different input because SQS tags are listed differently
func intersectSQSTags(tags map[string]*string, desired map[string]string) map[string]string {
	if tags == nil {
		return nil
	}
	actual := make(map[string]string)
	for k, v := range tags {
		vv := aws.StringValue(v)

		if _, found := desired[k]; found {
			actual[k] = vv
		}
	}
	if len(actual) == 0 && desired == nil {
		// Avoid problems with comparison between nil & {}
		return nil
	}
	return actual
}
<|endoftext|>


<commit_before>package sse

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

// Tests are accessing subscriber in a non-threadsafe way.
// Maybe fix this in the future so we can test with -race enabled

func TestStream(t *testing.T) {
	Convey("Given a new stream", t, func() {
		// New Stream
		s := newStream(1024)
		s.run()

		Convey("When adding a subscriber", func() {
			sub := s.addSubscriber()

			Convey("It should be stored", func() {
				So(len(s.subscribers), ShouldEqual, 1)
			})
			Convey("It should receive messages", func() {
				s.event <- []byte("test")
				msg, err := wait(sub.connection, time.Second*1)

				So(err, ShouldBeNil)
				So(string(msg), ShouldEqual, "test")
			})
		})

		Convey("When removing a subscriber", func() {
			s.addSubscriber()
			s.removeSubscriber(0)
			Convey("It should be removed from the list of subscribers", func() {
				So(len(s.subscribers), ShouldEqual, 0)
			})
		})

		Convey("When closing a subscriber down gracefully", func() {
			sub := s.addSubscriber()
			sub.close()
			time.Sleep(time.Millisecond * 100)
			Convey("It should be removed from the list of subscribers", func() {
				So(len(s.subscribers), ShouldEqual, 0)
			})
		})

		Convey("When adding multiple subscribers", func() {
			var subs []*Subscriber
			for i := 0; i < 10; i++ {
				subs = append(subs, s.addSubscriber())
			}

			// Wait for all subscribers to be added
			time.Sleep(time.Millisecond * 100)

			Convey("They should all receive messages", func() {
				s.event <- []byte("test")
				for _, sub := range subs {
					msg, err := wait(sub.connection, time.Second*1)
					So(err, ShouldBeNil)
					So(string(msg), ShouldEqual, "test")
				}
			})

			Convey("They should all shutdown gracefully when the stream is closed", func() {
				s.close()

				// Wait for all subscribers to close
				time.Sleep(time.Millisecond * 100)

				So(len(s.subscribers), ShouldEqual, 0)
			})

		})
	})

}
<commit_msg>added wait for removing subscriber on test<commit_after>package sse

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

// Tests are accessing subscriber in a non-threadsafe way.
// Maybe fix this in the future so we can test with -race enabled

func TestStream(t *testing.T) {
	Convey("Given a new stream", t, func() {
		// New Stream
		s := newStream(1024)
		s.run()

		Convey("When adding a subscriber", func() {
			sub := s.addSubscriber()

			Convey("It should be stored", func() {
				So(len(s.subscribers), ShouldEqual, 1)
			})
			Convey("It should receive messages", func() {
				s.event <- []byte("test")
				msg, err := wait(sub.connection, time.Second*1)

				So(err, ShouldBeNil)
				So(string(msg), ShouldEqual, "test")
			})
		})

		Convey("When removing a subscriber", func() {
			s.addSubscriber()
			time.Sleep(time.Millisecond * 100)
			s.removeSubscriber(0)

			Convey("It should be removed from the list of subscribers", func() {
				So(len(s.subscribers), ShouldEqual, 0)
			})
		})

		Convey("When closing a subscriber down gracefully", func() {
			sub := s.addSubscriber()
			sub.close()
			time.Sleep(time.Millisecond * 100)

			Convey("It should be removed from the list of subscribers", func() {
				So(len(s.subscribers), ShouldEqual, 0)
			})
		})

		Convey("When adding multiple subscribers", func() {
			var subs []*Subscriber
			for i := 0; i < 10; i++ {
				subs = append(subs, s.addSubscriber())
			}

			// Wait for all subscribers to be added
			time.Sleep(time.Millisecond * 100)

			Convey("They should all receive messages", func() {
				s.event <- []byte("test")
				for _, sub := range subs {
					msg, err := wait(sub.connection, time.Second*1)
					So(err, ShouldBeNil)
					So(string(msg), ShouldEqual, "test")
				}
			})

			Convey("They should all shutdown gracefully when the stream is closed", func() {
				s.close()

				// Wait for all subscribers to close
				time.Sleep(time.Millisecond * 100)

				So(len(s.subscribers), ShouldEqual, 0)
			})

		})
	})

}
<|endoftext|>