<commit_before>package cache

import (
	"context"
	"encoding/json"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/registry"
)

const refreshWhenExpiryWithin = time.Minute
const askForNewImagesInterval = time.Minute

// Warmer refreshes the information kept in the cache from remote
// registries.
type Warmer struct {
	clientFactory registry.ClientFactory
	cache         Client
	burst         int
	Priority      chan image.Name
	Notify        func()
}

// NewWarmer creates cache warmer that (when Loop is invoked) will
// periodically refresh the values kept in the cache.
func NewWarmer(cf registry.ClientFactory, cacheClient Client, burst int) (*Warmer, error) {
	if cf == nil || cacheClient == nil || burst <= 0 {
		return nil, errors.New("arguments must be non-nil (or > 0 in the case of burst)")
	}
	return &Warmer{
		clientFactory: cf,
		cache:         cacheClient,
		burst:         burst,
	}, nil
}

// .. and this is what we keep in the backlog
type backlogItem struct {
	image.Name
	registry.Credentials
}

// Loop continuously gets the images to populate the cache with,
// and populate the cache with them.
func (w *Warmer) Loop(logger log.Logger, stop <-chan struct{}, wg *sync.WaitGroup, imagesToFetchFunc func() registry.ImageCreds) {
	defer wg.Done()

	refresh := time.Tick(askForNewImagesInterval)
	imageCreds := imagesToFetchFunc()
	backlog := imageCredsToBacklog(imageCreds)

	// We have some fine control over how long to spend on each fetch
	// operation, since they are given a `context`. For now though,
	// just rattle through them one by one, however long they take.
	ctx := context.Background()

	// This loop acts keeps a kind of priority queue, whereby image
	// names coming in on the `Priority` channel are looked up first.
	// If there are none, images used in the cluster are refreshed;
	// but no more often than once every `askForNewImagesInterval`,
	// since there is no effective back-pressure on cache refreshes
	// and it would spin freely otherwise).
	for {
		select {
		case <-stop:
			logger.Log("stopping", "true")
			return
		case name := <-w.Priority:
			logger.Log("priority", name.String())
			// NB the implicit contract here is that the prioritised
			// image has to have been running the last time we
			// requested the credentials.
			if creds, ok := imageCreds[name]; ok {
				w.warm(ctx, logger, name, creds)
			} else {
				logger.Log("priority", name.String(), "err", "no creds available")
			}
			continue
		default:
		}

		if len(backlog) > 0 {
			im := backlog[0]
			backlog = backlog[1:]
			w.warm(ctx, logger, im.Name, im.Credentials)
		} else {
			select {
			case <-refresh:
				imageCreds = imagesToFetchFunc()
				backlog = imageCredsToBacklog(imageCreds)
			default:
			}
		}
	}
}

func imageCredsToBacklog(imageCreds registry.ImageCreds) []backlogItem {
	backlog := make([]backlogItem, len(imageCreds))
	var i int
	for name, cred := range imageCreds {
		backlog[i] = backlogItem{name, cred}
		i++
	}
	return backlog
}

func (w *Warmer) warm(ctx context.Context, logger log.Logger, id image.Name, creds registry.Credentials) {
	errorLogger := log.With(logger, "canonical_name", id.CanonicalName(), "auth", creds)
	client, err := w.clientFactory.ClientFor(id.CanonicalName(), creds)
	if err != nil {
		errorLogger.Log("err", err.Error())
		return
	}

	// This is what we're going to write back to the cache
	var repo ImageRepository
	repoKey := NewRepositoryKey(id.CanonicalName())
	bytes, _, err := w.cache.GetKey(repoKey)
	if err == nil {
		err = json.Unmarshal(bytes, &repo)
	} else if err == ErrNotCached {
		err = nil
	}

	if err != nil {
		errorLogger.Log("err", errors.Wrap(err, "fetching previous result from cache"))
		return
	}
	// Save for comparison later
	oldImages := repo.Images

	// Now we have the previous result; everything after will be
	// attempting to refresh that value. Whatever happens, at the end
	// we'll write something back.
	defer func() {
		bytes, err := json.Marshal(repo)
		if err == nil {
			err = w.cache.SetKey(repoKey, bytes)
		}
		if err != nil {
			errorLogger.Log("err", errors.Wrap(err, "writing result to cache"))
		}
	}()

	tags, err := client.Tags(ctx)
	if err != nil {
		if !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) && !strings.Contains(err.Error(), "net/http: request canceled") {
			errorLogger.Log("err", errors.Wrap(err, "requesting tags"))
			repo.LastError = err.Error()
		}
		return
	}

	newImages := map[string]image.Info{}

	// Create a list of manifests that need updating
	var toUpdate []image.Ref
	var missing, expired int
	for _, tag := range tags {
		// See if we have the manifest already cached
		newID := id.ToRef(tag)
		key := NewManifestKey(newID.CanonicalRef())
		bytes, expiry, err := w.cache.GetKey(key)
		// If err, then we don't have it yet. Update.
		switch {
		case err != nil:
			missing++
		case time.Until(expiry) < refreshWhenExpiryWithin:
			expired++
		default:
			var image image.Info
			if err := json.Unmarshal(bytes, &image); err == nil {
				newImages[tag] = image
				continue
			}
			missing++
		}
		toUpdate = append(toUpdate, newID)
	}

	var successCount int

	if len(toUpdate) > 0 {
		logger.Log("fetching", id.String(), "total", len(toUpdate), "expired", expired, "missing", missing)
		var successMx sync.Mutex

		// The upper bound for concurrent fetches against a single host is
		// w.Burst, so limit the number of fetching goroutines to that.
		fetchers := make(chan struct{}, w.burst)
		awaitFetchers := &sync.WaitGroup{}
	updates:
		for _, imID := range toUpdate {
			select {
			case <-ctx.Done():
				break updates
			case fetchers <- struct{}{}:
			}

			awaitFetchers.Add(1)
			go func(imageID image.Ref) {
				defer func() { awaitFetchers.Done(); <-fetchers }()
				// Get the image from the remote
				img, err := client.Manifest(ctx, imageID.Tag)
				if err != nil {
					if err, ok := errors.Cause(err).(net.Error); ok && err.Timeout() {
						// This was due to a context timeout, don't bother logging
						return
					}
					errorLogger.Log("err", errors.Wrap(err, "requesting manifests"))
					return
				}

				key := NewManifestKey(img.ID.CanonicalRef())
				// Write back to memcached
				val, err := json.Marshal(img)
				if err != nil {
					errorLogger.Log("err", errors.Wrap(err, "serializing tag to store in cache"))
					return
				}
				err = w.cache.SetKey(key, val)
				if err != nil {
					errorLogger.Log("err", errors.Wrap(err, "storing manifests in cache"))
					return
				}
				successMx.Lock()
				successCount++
				newImages[imageID.Tag] = img
				successMx.Unlock()
			}(imID)
		}
		awaitFetchers.Wait()
		logger.Log("updated", id.String(), "count", successCount)
	}

	// We managed to fetch new metadata for everything we were missing
	// (if anything). Ratchet the result forward.
	if successCount == len(toUpdate) {
		repo = ImageRepository{
			LastUpdate: time.Now(),
			Images:     newImages,
		}
	}

	if w.Notify != nil {
		cacheTags := StringSet{}
		for t := range oldImages {
			cacheTags[t] = struct{}{}
		}

		// If there's more tags than there used to be, there must be
		// at least one new tag.
		if len(cacheTags) < len(tags) {
			w.Notify()
			return
		}
		// Otherwise, check whether there are any entries in the
		// fetched tags that aren't in the cached tags.
		tagSet := NewStringSet(tags)
		if !tagSet.Subset(cacheTags) {
			w.Notify()
		}
	}
}

// StringSet is a set of strings.
type StringSet map[string]struct{}

// NewStringSet returns a StringSet containing exactly the strings
// given as arguments.
func NewStringSet(ss []string) StringSet {
	res := StringSet{}
	for _, s := range ss {
		res[s] = struct{}{}
	}
	return res
}

// Subset returns true if `s` is a subset of `t` (including the case
// of having the same members).
func (s StringSet) Subset(t StringSet) bool {
	for k := range s {
		if _, ok := t[k]; !ok {
			return false
		}
	}
	return true
}
<commit_msg>Avoid spin in cache refresh loop To implement the loop as described in the comment above it, when the backlog is empty we need to block until the refresh tick fires (or a priority image comes in). By having a default case, it's possible to spin.<commit_after>package cache

import (
	"context"
	"encoding/json"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/registry"
)

const refreshWhenExpiryWithin = time.Minute
const askForNewImagesInterval = time.Minute

// Warmer refreshes the information kept in the cache from remote
// registries.
type Warmer struct {
	clientFactory registry.ClientFactory
	cache         Client
	burst         int
	Priority      chan image.Name
	Notify        func()
}

// NewWarmer creates cache warmer that (when Loop is invoked) will
// periodically refresh the values kept in the cache.
func NewWarmer(cf registry.ClientFactory, cacheClient Client, burst int) (*Warmer, error) {
	if cf == nil || cacheClient == nil || burst <= 0 {
		return nil, errors.New("arguments must be non-nil (or > 0 in the case of burst)")
	}
	return &Warmer{
		clientFactory: cf,
		cache:         cacheClient,
		burst:         burst,
	}, nil
}

// .. and this is what we keep in the backlog
type backlogItem struct {
	image.Name
	registry.Credentials
}

// Loop continuously gets the images to populate the cache with,
// and populate the cache with them.
func (w *Warmer) Loop(logger log.Logger, stop <-chan struct{}, wg *sync.WaitGroup, imagesToFetchFunc func() registry.ImageCreds) {
	defer wg.Done()

	refresh := time.Tick(askForNewImagesInterval)
	imageCreds := imagesToFetchFunc()
	backlog := imageCredsToBacklog(imageCreds)

	// We have some fine control over how long to spend on each fetch
	// operation, since they are given a `context`. For now though,
	// just rattle through them one by one, however long they take.
	ctx := context.Background()

	// NB the implicit contract here is that the prioritised
	// image has to have been running the last time we
	// requested the credentials.
	priorityWarm := func (name image.Name) {
		logger.Log("priority", name.String())
		if creds, ok := imageCreds[name]; ok {
			w.warm(ctx, logger, name, creds)
		} else {
			logger.Log("priority", name.String(), "err", "no creds available")
		}
	}

	// This loop acts keeps a kind of priority queue, whereby image
	// names coming in on the `Priority` channel are looked up first.
	// If there are none, images used in the cluster are refreshed;
	// but no more often than once every `askForNewImagesInterval`,
	// since there is no effective back-pressure on cache refreshes
	// and it would spin freely otherwise).
	for {
		select {
		case <-stop:
			logger.Log("stopping", "true")
			return
		case name := <-w.Priority:
			priorityWarm(name)
			continue
		default:
		}

		if len(backlog) > 0 {
			im := backlog[0]
			backlog = backlog[1:]
			w.warm(ctx, logger, im.Name, im.Credentials)
		} else {
			select {
			case <-stop:
				logger.Log("stopping", "true")
				return
			case <-refresh:
				imageCreds = imagesToFetchFunc()
				backlog = imageCredsToBacklog(imageCreds)
			case name := <-w.Priority:
				priorityWarm(name)
			}
		}
	}
}

func imageCredsToBacklog(imageCreds registry.ImageCreds) []backlogItem {
	backlog := make([]backlogItem, len(imageCreds))
	var i int
	for name, cred := range imageCreds {
		backlog[i] = backlogItem{name, cred}
		i++
	}
	return backlog
}

func (w *Warmer) warm(ctx context.Context, logger log.Logger, id image.Name, creds registry.Credentials) {
	errorLogger := log.With(logger, "canonical_name", id.CanonicalName(), "auth", creds)
	client, err := w.clientFactory.ClientFor(id.CanonicalName(), creds)
	if err != nil {
		errorLogger.Log("err", err.Error())
		return
	}

	// This is what we're going to write back to the cache
	var repo ImageRepository
	repoKey := NewRepositoryKey(id.CanonicalName())
	bytes, _, err := w.cache.GetKey(repoKey)
	if err == nil {
		err = json.Unmarshal(bytes, &repo)
	} else if err == ErrNotCached {
		err = nil
	}

	if err != nil {
		errorLogger.Log("err", errors.Wrap(err, "fetching previous result from cache"))
		return
	}
	// Save for comparison later
	oldImages := repo.Images

	// Now we have the previous result; everything after will be
	// attempting to refresh that value. Whatever happens, at the end
	// we'll write something back.
	defer func() {
		bytes, err := json.Marshal(repo)
		if err == nil {
			err = w.cache.SetKey(repoKey, bytes)
		}
		if err != nil {
			errorLogger.Log("err", errors.Wrap(err, "writing result to cache"))
		}
	}()

	tags, err := client.Tags(ctx)
	if err != nil {
		if !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) && !strings.Contains(err.Error(), "net/http: request canceled") {
			errorLogger.Log("err", errors.Wrap(err, "requesting tags"))
			repo.LastError = err.Error()
		}
		return
	}

	newImages := map[string]image.Info{}

	// Create a list of manifests that need updating
	var toUpdate []image.Ref
	var missing, expired int
	for _, tag := range tags {
		// See if we have the manifest already cached
		newID := id.ToRef(tag)
		key := NewManifestKey(newID.CanonicalRef())
		bytes, expiry, err := w.cache.GetKey(key)
		// If err, then we don't have it yet. Update.
		switch {
		case err != nil:
			missing++
		case time.Until(expiry) < refreshWhenExpiryWithin:
			expired++
		default:
			var image image.Info
			if err := json.Unmarshal(bytes, &image); err == nil {
				newImages[tag] = image
				continue
			}
			missing++
		}
		toUpdate = append(toUpdate, newID)
	}

	var successCount int

	if len(toUpdate) > 0 {
		logger.Log("fetching", id.String(), "total", len(toUpdate), "expired", expired, "missing", missing)
		var successMx sync.Mutex

		// The upper bound for concurrent fetches against a single host is
		// w.Burst, so limit the number of fetching goroutines to that.
		fetchers := make(chan struct{}, w.burst)
		awaitFetchers := &sync.WaitGroup{}
	updates:
		for _, imID := range toUpdate {
			select {
			case <-ctx.Done():
				break updates
			case fetchers <- struct{}{}:
			}

			awaitFetchers.Add(1)
			go func(imageID image.Ref) {
				defer func() { awaitFetchers.Done(); <-fetchers }()
				// Get the image from the remote
				img, err := client.Manifest(ctx, imageID.Tag)
				if err != nil {
					if err, ok := errors.Cause(err).(net.Error); ok && err.Timeout() {
						// This was due to a context timeout, don't bother logging
						return
					}
					errorLogger.Log("err", errors.Wrap(err, "requesting manifests"))
					return
				}

				key := NewManifestKey(img.ID.CanonicalRef())
				// Write back to memcached
				val, err := json.Marshal(img)
				if err != nil {
					errorLogger.Log("err", errors.Wrap(err, "serializing tag to store in cache"))
					return
				}
				err = w.cache.SetKey(key, val)
				if err != nil {
					errorLogger.Log("err", errors.Wrap(err, "storing manifests in cache"))
					return
				}
				successMx.Lock()
				successCount++
				newImages[imageID.Tag] = img
				successMx.Unlock()
			}(imID)
		}
		awaitFetchers.Wait()
		logger.Log("updated", id.String(), "count", successCount)
	}

	// We managed to fetch new metadata for everything we were missing
	// (if anything). Ratchet the result forward.
	if successCount == len(toUpdate) {
		repo = ImageRepository{
			LastUpdate: time.Now(),
			Images:     newImages,
		}
	}

	if w.Notify != nil {
		cacheTags := StringSet{}
		for t := range oldImages {
			cacheTags[t] = struct{}{}
		}

		// If there's more tags than there used to be, there must be
		// at least one new tag.
		if len(cacheTags) < len(tags) {
			w.Notify()
			return
		}
		// Otherwise, check whether there are any entries in the
		// fetched tags that aren't in the cached tags.
		tagSet := NewStringSet(tags)
		if !tagSet.Subset(cacheTags) {
			w.Notify()
		}
	}
}

// StringSet is a set of strings.
type StringSet map[string]struct{}

// NewStringSet returns a StringSet containing exactly the strings
// given as arguments.
func NewStringSet(ss []string) StringSet {
	res := StringSet{}
	for _, s := range ss {
		res[s] = struct{}{}
	}
	return res
}

// Subset returns true if `s` is a subset of `t` (including the case
// of having the same members).
func (s StringSet) Subset(t StringSet) bool {
	for k := range s {
		if _, ok := t[k]; !ok {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package repository

import (
	"errors"
	"fmt"

	"github.com/hoffie/larasync/repository/odf"
)

// Transaction represents a server side transaction for specific NIBs
// which is used to synchronize the different clients.
type Transaction struct {
	ID         int64
	NIBIDs     []string
	PreviousID int64
}

// newTransactionFromPb returns a new Transaction from the
// protobuf Transaction.
func newTransactionFromPb(pbTransaction *odf.Transaction) *Transaction {
	return &Transaction{
		ID:         pbTransaction.GetID(),
		PreviousID: pbTransaction.GetPreviousID(),
		NIBIDs:     pbTransaction.GetNIBIDs(),
	}
}

// toPb converts this Transaction to a protobuf Transaction.
// This is used by the encoder.
func (t *Transaction) toPb() (*odf.Transaction, error) {
	if t.ID == 0 {
		return nil, errors.New("Transaction ID must not be empty")
	}
	if len(t.NIBIDs) == 0 {
		return nil, fmt.Errorf(
			"The transaction with ID %d has no NIB IDs",
			t.ID,
		)
	}
	protoTransaction := &odf.Transaction{
		ID:         &t.ID,
		PreviousID: nil,
		NIBIDs:     t.NIBIDs}
	if t.PreviousID != 0 {
		protoTransaction.PreviousID = &t.PreviousID
	}
	return protoTransaction, nil
}

// nibUUIDsFromTransactions returns all uuids from a list of transactions.
func nibUUIDsFromTransactions(transactions []*Transaction) <-chan string {
	nibUUIDChannel := make(chan string, 100)
	go func() {
		for _, transaction := range transactions {
			for _, nibID := range transaction.NIBIDs {
				nibUUIDChannel <- nibID
			}
		}
		close(nibUUIDChannel)
	}()
	return nibUUIDChannel
}
<commit_msg>repository: Implemented the IDString() function in the Transaction object.<commit_after>package repository

import (
	"errors"
	"fmt"
	"strconv"

	"github.com/hoffie/larasync/repository/odf"
)

// Transaction represents a server side transaction for specific NIBs
// which is used to synchronize the different clients.
type Transaction struct {
	ID         int64
	NIBIDs     []string
	PreviousID int64
}

// newTransactionFromPb returns a new Transaction from the
// protobuf Transaction.
func newTransactionFromPb(pbTransaction *odf.Transaction) *Transaction {
	return &Transaction{
		ID:         pbTransaction.GetID(),
		PreviousID: pbTransaction.GetPreviousID(),
		NIBIDs:     pbTransaction.GetNIBIDs(),
	}
}

// toPb converts this Transaction to a protobuf Transaction.
// This is used by the encoder.
func (t *Transaction) toPb() (*odf.Transaction, error) {
	if t.ID == 0 {
		return nil, errors.New("Transaction ID must not be empty")
	}
	if len(t.NIBIDs) == 0 {
		return nil, fmt.Errorf(
			"The transaction with ID %d has no NIB IDs",
			t.ID,
		)
	}
	protoTransaction := &odf.Transaction{
		ID:         &t.ID,
		PreviousID: nil,
		NIBIDs:     t.NIBIDs}
	if t.PreviousID != 0 {
		protoTransaction.PreviousID = &t.PreviousID
	}
	return protoTransaction, nil
}

// IDString returns the ID of this Transaction as a string.
func (t *Transaction) IDString() string {
	return strconv.FormatInt(t.ID, 10)
}

// nibUUIDsFromTransactions returns all uuids from a list of transactions.
func nibUUIDsFromTransactions(transactions []*Transaction) <-chan string {
	nibUUIDChannel := make(chan string, 100)
	go func() {
		for _, transaction := range transactions {
			for _, nibID := range transaction.NIBIDs {
				nibUUIDChannel <- nibID
			}
		}
		close(nibUUIDChannel)
	}()
	return nibUUIDChannel
}
<|endoftext|>


<commit_before>package admin

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
	"github.com/qor/action_bar"
	"github.com/qor/activity"
	"github.com/qor/admin"
	"github.com/qor/i18n/exchange_actions"
	"github.com/qor/media_library"
	"github.com/qor/qor"
	"github.com/qor/qor-example/app/models"
	"github.com/qor/qor-example/config/admin/bindatafs"
	"github.com/qor/qor-example/config/auth"
	"github.com/qor/qor-example/config/i18n"
	"github.com/qor/qor-example/db"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/utils"
	"github.com/qor/transition"
	"github.com/qor/validations"
)

var Admin *admin.Admin
var ActionBar *action_bar.ActionBar
var Countries = []string{"China", "Japan", "USA"}

func init() {
	Admin = admin.New(&qor.Config{DB: db.Publish.DraftDB()})
	Admin.SetSiteName("Qor DEMO")
	Admin.SetAuth(auth.AdminAuth{})
	Admin.SetAssetFS(bindatafs.AssetFS)

	// Add Dashboard
	Admin.AddMenu(&admin.Menu{Name: "Dashboard", Link: "/admin"})

	// Add Asset Manager, for rich editor
	assetManager := Admin.AddResource(&media_library.AssetManager{}, &admin.Config{Invisible: true})

	// Add Product
	product := Admin.AddResource(&models.Product{}, &admin.Config{Menu: []string{"Product Management"}})
	product.Meta(&admin.Meta{Name: "MadeCountry", Type: "select_one", Collection: Countries})
	product.Meta(&admin.Meta{Name: "Description", Type: "rich_editor", Resource: assetManager})

	colorVariationMeta := product.Meta(&admin.Meta{Name: "ColorVariations"})
	colorVariation := colorVariationMeta.Resource
	colorVariation.NewAttrs("-Product")
	colorVariation.EditAttrs("-Product")

	sizeVariationMeta := colorVariation.Meta(&admin.Meta{Name: "SizeVariations"})
	sizeVariation := sizeVariationMeta.Resource
	sizeVariation.NewAttrs("-ColorVariation")
	sizeVariation.EditAttrs(
		&admin.Section{
			Rows: [][]string{
				{"Size", "AvailableQuantity"},
			},
		},
	)

	product.SearchAttrs("Name", "Code", "Category.Name", "Brand.Name")
	product.EditAttrs(
		&admin.Section{
			Title: "Basic Information",
			Rows: [][]string{
				{"Name"},
				{"Code", "Price"},
				{"Enabled"},
			}},
		&admin.Section{
			Title: "Organization",
			Rows: [][]string{
				{"Category", "MadeCountry"},
				{"Collections"},
			}},
		"Description",
		"ColorVariations",
	)

	for _, country := range Countries {
		var country = country
		product.Scope(&admin.Scope{Name: country, Group: "Made Country", Handle: func(db *gorm.DB, ctx *qor.Context) *gorm.DB {
			return db.Where("made_country = ?", country)
		}})
	}

	product.IndexAttrs("-ColorVariations")

	product.Action(&admin.Action{
		Name: "View On Site",
		URL: func(record interface{}, context *admin.Context) string {
			if product, ok := record.(*models.Product); ok {
				return fmt.Sprintf("/products/%v", product.Code)
			}
			return "#"
		},
		Modes: []string{"menu_item", "edit"},
	})

	product.Action(&admin.Action{
		Name: "Disable",
		Handle: func(arg *admin.ActionArgument) error {
			for _, record := range arg.FindSelectedRecords() {
				arg.Context.DB.Model(record.(*models.Product)).Update("enabled", false)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if product, ok := record.(*models.Product); ok {
				return product.Enabled == true
			}
			return true
		},
		Modes: []string{"index", "edit", "menu_item"},
	})

	product.Action(&admin.Action{
		Name: "Enable",
		Handle: func(arg *admin.ActionArgument) error {
			for _, record := range arg.FindSelectedRecords() {
				arg.Context.DB.Model(record.(*models.Product)).Update("enabled", true)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if product, ok := record.(*models.Product); ok {
				return product.Enabled == false
			}
			return true
		},
		Modes: []string{"index", "edit", "menu_item"},
	})

	Admin.AddResource(&models.Color{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Size{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Category{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Collection{}, &admin.Config{Menu: []string{"Product Management"}})

	// Add Order
	order := Admin.AddResource(&models.Order{}, &admin.Config{Menu: []string{"Order Management"}})
	order.Meta(&admin.Meta{Name: "ShippingAddress", Type: "single_edit"})
	order.Meta(&admin.Meta{Name: "BillingAddress", Type: "single_edit"})
	order.Meta(&admin.Meta{Name: "ShippedAt", Type: "date"})

	orderItemMeta := order.Meta(&admin.Meta{Name: "OrderItems"})
	orderItemMeta.Resource.Meta(&admin.Meta{Name: "SizeVariation", Type: "select_one", Collection: sizeVariationCollection})
	orderItemMeta.Resource.NewAttrs("-State")
	orderItemMeta.Resource.EditAttrs("-State")

	// define scopes for Order
	for _, state := range []string{"checkout", "cancelled", "paid", "paid_cancelled", "processing", "shipped", "returned"} {
		var state = state
		order.Scope(&admin.Scope{
			Name:  state,
			Label: strings.Title(strings.Replace(state, "_", " ", -1)),
			Group: "Order Status",
			Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
				return db.Where(models.Order{Transition: transition.Transition{State: state}})
			},
		})
	}

	// define actions for Order
	type trackingNumberArgument struct {
		TrackingNumber string
	}

	order.Action(&admin.Action{
		Name: "Processing",
		Handle: func(argument *admin.ActionArgument) error {
			for _, order := range argument.FindSelectedRecords() {
				db := argument.Context.GetDB()
				if err := models.OrderState.Trigger("process", order.(*models.Order), db); err != nil {
					return err
				}
				db.Select("state").Save(order)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				return order.State == "paid"
			}
			return false
		},
		Modes: []string{"show", "menu_item"},
	})
	order.Action(&admin.Action{
		Name: "Ship",
		Handle: func(argument *admin.ActionArgument) error {
			var (
				tx                     = argument.Context.GetDB().Begin()
				trackingNumberArgument = argument.Argument.(*trackingNumberArgument)
			)

			if trackingNumberArgument.TrackingNumber != "" {
				for _, record := range argument.FindSelectedRecords() {
					order := record.(*models.Order)
					order.TrackingNumber = &trackingNumberArgument.TrackingNumber
					models.OrderState.Trigger("ship", order, tx, "tracking number "+trackingNumberArgument.TrackingNumber)
					if err := tx.Save(order).Error; err != nil {
						tx.Rollback()
						return err
					}
				}
			} else {
				return errors.New("invalid shipment number")
			}

			tx.Commit()
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				return order.State == "processing"
			}
			return false
		},
		Resource: Admin.NewResource(&trackingNumberArgument{}),
		Modes:    []string{"show", "menu_item"},
	})

	order.Action(&admin.Action{
		Name: "Cancel",
		Handle: func(argument *admin.ActionArgument) error {
			for _, order := range argument.FindSelectedRecords() {
				db := argument.Context.GetDB()
				if err := models.OrderState.Trigger("cancel", order.(*models.Order), db); err != nil {
					return err
				}
				db.Select("state").Save(order)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				for _, state := range []string{"draft", "checkout", "paid", "processing"} {
					if order.State == state {
						return true
					}
				}
			}
			return false
		},
		Modes: []string{"index", "show", "menu_item"},
	})

	order.IndexAttrs("User", "PaymentAmount", "ShippedAt", "CancelledAt", "State", "ShippingAddress")
	order.NewAttrs("-DiscountValue", "-AbandonedReason", "-CancelledAt")
	order.EditAttrs("-DiscountValue", "-AbandonedReason", "-CancelledAt", "-State")
	order.ShowAttrs("-DiscountValue", "-State")
	order.SearchAttrs("User.Name", "User.Email", "ShippingAddress.ContactName", "ShippingAddress.Address1", "ShippingAddress.Address2")

	// Add activity for order
	activity.Register(order)

	// Define another resource for same model
	abandonedOrder := Admin.AddResource(&models.Order{}, &admin.Config{Name: "Abandoned Order", Menu: []string{"Order Management"}})
	abandonedOrder.Meta(&admin.Meta{Name: "ShippingAddress", Type: "single_edit"})
	abandonedOrder.Meta(&admin.Meta{Name: "BillingAddress", Type: "single_edit"})

	// Define default scope for abandoned orders
	abandonedOrder.Scope(&admin.Scope{
		Default: true,
		Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
			return db.Where("abandoned_reason IS NOT NULL AND abandoned_reason <> ?", "")
		},
	})

	// Define scopes for abandoned orders
	for _, amount := range []int{5000, 10000, 20000} {
		var amount = amount
		abandonedOrder.Scope(&admin.Scope{
			Name:  fmt.Sprint(amount),
			Group: "Amount Greater Than",
			Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
				return db.Where("payment_amount > ?", amount)
			},
		})
	}

	abandonedOrder.IndexAttrs("-ShippingAddress", "-BillingAddress", "-DiscountValue", "-OrderItems")
	abandonedOrder.NewAttrs("-DiscountValue")
	abandonedOrder.EditAttrs("-DiscountValue")
	abandonedOrder.ShowAttrs("-DiscountValue")

	// Add Store
	store := Admin.AddResource(&models.Store{}, &admin.Config{Menu: []string{"Store Management"}})
	store.AddValidator(func(record interface{}, metaValues *resource.MetaValues, context *qor.Context) error {
		if meta := metaValues.Get("Name"); meta != nil {
			if name := utils.ToString(meta.Value); strings.TrimSpace(name) == "" {
				return validations.NewError(record, "Name", "Name can't be blank")
			}
		}
		return nil
	})

	// Add Translations
	Admin.AddResource(i18n.I18n, &admin.Config{Menu: []string{"Site Management"}})

	// Add SEOSetting
	Admin.AddResource(&models.SEOSetting{}, &admin.Config{Menu: []string{"Site Management"}, Singleton: true})

	// Add Setting
	Admin.AddResource(&models.Setting{}, &admin.Config{Singleton: true})

	// Add User
	user := Admin.AddResource(&models.User{})
	user.Meta(&admin.Meta{Name: "Gender", Type: "select_one", Collection: []string{"Male", "Female", "Unknown"}})

	user.IndexAttrs("ID", "Email", "Name", "Gender", "Role")
	user.ShowAttrs(
		&admin.Section{
			Title: "Basic Information",
			Rows: [][]string{
				{"Name"},
				{"Email", "Password"},
				{"Gender", "Role"},
			}},
		"Addresses",
	)
	user.EditAttrs(user.ShowAttrs())

	// Add Worker
	Worker := getWorker()
	Admin.AddResource(Worker)

	db.Publish.SetWorker(Worker)
	exchange_actions.RegisterExchangeJobs(i18n.I18n, Worker)

	// Add Publish
	Admin.AddResource(db.Publish, &admin.Config{Singleton: true})

	// Add Search Center Resources
	Admin.AddSearchResource(product, user, order)

	// Add ActionBar
	ActionBar = action_bar.New(Admin, auth.AdminAuth{})
	ActionBar.RegisterAction(&action_bar.Action{Name: "Admin Dashboard", Link: "/admin"})

	initFuncMap()
	initRouter()
}

func sizeVariationCollection(resource interface{}, context *qor.Context) (results [][]string) {
	for _, sizeVariation := range models.SizeVariations() {
		results = append(results, []string{strconv.Itoa(int(sizeVariation.ID)), sizeVariation.Stringify()})
	}
	return
}
<commit_msg>Use publish draft mode for admin<commit_after>package admin

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
	"github.com/qor/action_bar"
	"github.com/qor/activity"
	"github.com/qor/admin"
	"github.com/qor/i18n/exchange_actions"
	"github.com/qor/media_library"
	"github.com/qor/qor"
	"github.com/qor/qor-example/app/models"
	"github.com/qor/qor-example/config/admin/bindatafs"
	"github.com/qor/qor-example/config/auth"
	"github.com/qor/qor-example/config/i18n"
	"github.com/qor/qor-example/db"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/utils"
	"github.com/qor/transition"
	"github.com/qor/validations"
)

var Admin *admin.Admin
var ActionBar *action_bar.ActionBar
var Countries = []string{"China", "Japan", "USA"}

func init() {
	Admin = admin.New(&qor.Config{DB: db.DB.Set("publish:draft_mode", true)})
	Admin.SetSiteName("Qor DEMO")
	Admin.SetAuth(auth.AdminAuth{})
	Admin.SetAssetFS(bindatafs.AssetFS)

	// Add Dashboard
	Admin.AddMenu(&admin.Menu{Name: "Dashboard", Link: "/admin"})

	// Add Asset Manager, for rich editor
	assetManager := Admin.AddResource(&media_library.AssetManager{}, &admin.Config{Invisible: true})

	// Add Product
	product := Admin.AddResource(&models.Product{}, &admin.Config{Menu: []string{"Product Management"}})
	product.Meta(&admin.Meta{Name: "MadeCountry", Type: "select_one", Collection: Countries})
	product.Meta(&admin.Meta{Name: "Description", Type: "rich_editor", Resource: assetManager})

	colorVariationMeta := product.Meta(&admin.Meta{Name: "ColorVariations"})
	colorVariation := colorVariationMeta.Resource
	colorVariation.NewAttrs("-Product")
	colorVariation.EditAttrs("-Product")

	sizeVariationMeta := colorVariation.Meta(&admin.Meta{Name: "SizeVariations"})
	sizeVariation := sizeVariationMeta.Resource
	sizeVariation.NewAttrs("-ColorVariation")
	sizeVariation.EditAttrs(
		&admin.Section{
			Rows: [][]string{
				{"Size", "AvailableQuantity"},
			},
		},
	)

	product.SearchAttrs("Name", "Code", "Category.Name", "Brand.Name")
	product.EditAttrs(
		&admin.Section{
			Title: "Basic Information",
			Rows: [][]string{
				{"Name"},
				{"Code", "Price"},
				{"Enabled"},
			}},
		&admin.Section{
			Title: "Organization",
			Rows: [][]string{
				{"Category", "MadeCountry"},
				{"Collections"},
			}},
		"Description",
		"ColorVariations",
	)

	for _, country := range Countries {
		var country = country
		product.Scope(&admin.Scope{Name: country, Group: "Made Country", Handle: func(db *gorm.DB, ctx *qor.Context) *gorm.DB {
			return db.Where("made_country = ?", country)
		}})
	}

	product.IndexAttrs("-ColorVariations")

	product.Action(&admin.Action{
		Name: "View On Site",
		URL: func(record interface{}, context *admin.Context) string {
			if product, ok := record.(*models.Product); ok {
				return fmt.Sprintf("/products/%v", product.Code)
			}
			return "#"
		},
		Modes: []string{"menu_item", "edit"},
	})

	product.Action(&admin.Action{
		Name: "Disable",
		Handle: func(arg *admin.ActionArgument) error {
			for _, record := range arg.FindSelectedRecords() {
				arg.Context.DB.Model(record.(*models.Product)).Update("enabled", false)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if product, ok := record.(*models.Product); ok {
				return product.Enabled == true
			}
			return true
		},
		Modes: []string{"index", "edit", "menu_item"},
	})

	product.Action(&admin.Action{
		Name: "Enable",
		Handle: func(arg *admin.ActionArgument) error {
			for _, record := range arg.FindSelectedRecords() {
				arg.Context.DB.Model(record.(*models.Product)).Update("enabled", true)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if product, ok := record.(*models.Product); ok {
				return product.Enabled == false
			}
			return true
		},
		Modes: []string{"index", "edit", "menu_item"},
	})

	Admin.AddResource(&models.Color{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Size{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Category{}, &admin.Config{Menu: []string{"Product Management"}})
	Admin.AddResource(&models.Collection{}, &admin.Config{Menu: []string{"Product Management"}})

	// Add Order
	order := Admin.AddResource(&models.Order{}, &admin.Config{Menu: []string{"Order Management"}})
	order.Meta(&admin.Meta{Name: "ShippingAddress", Type: "single_edit"})
	order.Meta(&admin.Meta{Name: "BillingAddress", Type: "single_edit"})
	order.Meta(&admin.Meta{Name: "ShippedAt", Type: "date"})

	orderItemMeta := order.Meta(&admin.Meta{Name: "OrderItems"})
	orderItemMeta.Resource.Meta(&admin.Meta{Name: "SizeVariation", Type: "select_one", Collection: sizeVariationCollection})
	orderItemMeta.Resource.NewAttrs("-State")
	orderItemMeta.Resource.EditAttrs("-State")

	// define scopes for Order
	for _, state := range []string{"checkout", "cancelled", "paid", "paid_cancelled", "processing", "shipped", "returned"} {
		var state = state
		order.Scope(&admin.Scope{
			Name:  state,
			Label: strings.Title(strings.Replace(state, "_", " ", -1)),
			Group: "Order Status",
			Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
				return db.Where(models.Order{Transition: transition.Transition{State: state}})
			},
		})
	}

	// define actions for Order
	type trackingNumberArgument struct {
		TrackingNumber string
	}

	order.Action(&admin.Action{
		Name: "Processing",
		Handle: func(argument *admin.ActionArgument) error {
			for _, order := range argument.FindSelectedRecords() {
				db := argument.Context.GetDB()
				if err := models.OrderState.Trigger("process", order.(*models.Order), db); err != nil {
					return err
				}
				db.Select("state").Save(order)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				return order.State == "paid"
			}
			return false
		},
		Modes: []string{"show", "menu_item"},
	})
	order.Action(&admin.Action{
		Name: "Ship",
		Handle: func(argument *admin.ActionArgument) error {
			var (
				tx                     = argument.Context.GetDB().Begin()
				trackingNumberArgument = argument.Argument.(*trackingNumberArgument)
			)

			if trackingNumberArgument.TrackingNumber != "" {
				for _, record := range argument.FindSelectedRecords() {
					order := record.(*models.Order)
					order.TrackingNumber = &trackingNumberArgument.TrackingNumber
					models.OrderState.Trigger("ship", order, tx, "tracking number "+trackingNumberArgument.TrackingNumber)
					if err := tx.Save(order).Error; err != nil {
						tx.Rollback()
						return err
					}
				}
			} else {
				return errors.New("invalid shipment number")
			}

			tx.Commit()
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				return order.State == "processing"
			}
			return false
		},
		Resource: Admin.NewResource(&trackingNumberArgument{}),
		Modes:    []string{"show", "menu_item"},
	})

	order.Action(&admin.Action{
		Name: "Cancel",
		Handle: func(argument *admin.ActionArgument) error {
			for _, order := range argument.FindSelectedRecords() {
				db := argument.Context.GetDB()
				if err := models.OrderState.Trigger("cancel", order.(*models.Order), db); err != nil {
					return err
				}
				db.Select("state").Save(order)
			}
			return nil
		},
		Visible: func(record interface{}, context *admin.Context) bool {
			if order, ok := record.(*models.Order); ok {
				for _, state := range []string{"draft", "checkout", "paid", "processing"} {
					if order.State == state {
						return true
					}
				}
			}
			return false
		},
		Modes: []string{"index", "show", "menu_item"},
	})

	order.IndexAttrs("User", "PaymentAmount", "ShippedAt", "CancelledAt", "State", "ShippingAddress")
	order.NewAttrs("-DiscountValue", "-AbandonedReason", "-CancelledAt")
	order.EditAttrs("-DiscountValue", "-AbandonedReason", "-CancelledAt", "-State")
	order.ShowAttrs("-DiscountValue", "-State")
	order.SearchAttrs("User.Name", "User.Email", "ShippingAddress.ContactName", "ShippingAddress.Address1", "ShippingAddress.Address2")

	// Add activity for order
	activity.Register(order)

	// Define another resource for same model
	abandonedOrder := Admin.AddResource(&models.Order{}, &admin.Config{Name: "Abandoned Order", Menu: []string{"Order Management"}})
	abandonedOrder.Meta(&admin.Meta{Name: "ShippingAddress", Type: "single_edit"})
	abandonedOrder.Meta(&admin.Meta{Name: "BillingAddress", Type: "single_edit"})

	// Define default scope for abandoned orders
	abandonedOrder.Scope(&admin.Scope{
		Default: true,
		Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
			return db.Where("abandoned_reason IS NOT NULL AND abandoned_reason <> ?", "")
		},
	})

	// Define scopes for abandoned orders
	for _, amount := range []int{5000, 10000, 20000} {
		var amount = amount
		abandonedOrder.Scope(&admin.Scope{
			Name:  fmt.Sprint(amount),
			Group: "Amount Greater Than",
			Handle: func(db *gorm.DB, context *qor.Context) *gorm.DB {
				return db.Where("payment_amount > ?", amount)
			},
		})
	}

	abandonedOrder.IndexAttrs("-ShippingAddress", "-BillingAddress", "-DiscountValue", "-OrderItems")
	abandonedOrder.NewAttrs("-DiscountValue")
	abandonedOrder.EditAttrs("-DiscountValue")
	abandonedOrder.ShowAttrs("-DiscountValue")

	// Add Store
	store := Admin.AddResource(&models.Store{}, &admin.Config{Menu: []string{"Store Management"}})
	store.AddValidator(func(record interface{}, metaValues *resource.MetaValues, context *qor.Context) error {
		if meta := metaValues.Get("Name"); meta != nil {
			if name := utils.ToString(meta.Value); strings.TrimSpace(name) == "" {
				return validations.NewError(record, "Name", "Name can't be blank")
			}
		}
		return nil
	})

	// Add Translations
	Admin.AddResource(i18n.I18n, &admin.Config{Menu: []string{"Site Management"}})

	// Add SEOSetting
	Admin.AddResource(&models.SEOSetting{}, &admin.Config{Menu: []string{"Site Management"}, Singleton: true})

	// Add Setting
	Admin.AddResource(&models.Setting{}, &admin.Config{Singleton: true})

	// Add User
	user := Admin.AddResource(&models.User{})
	user.Meta(&admin.Meta{Name: "Gender", Type: "select_one", Collection: []string{"Male", "Female", "Unknown"}})

	user.IndexAttrs("ID", "Email", "Name", "Gender", "Role")
	user.ShowAttrs(
		&admin.Section{
			Title: "Basic Information",
			Rows: [][]string{
				{"Name"},
				{"Email", "Password"},
				{"Gender", "Role"},
			}},
		"Addresses",
	)
	user.EditAttrs(user.ShowAttrs())

	// Add Worker
	Worker := getWorker()
	Admin.AddResource(Worker)

	db.Publish.SetWorker(Worker)
	exchange_actions.RegisterExchangeJobs(i18n.I18n, Worker)

	// Add Publish
	Admin.AddResource(db.Publish, &admin.Config{Singleton: true})

	// Add Search Center Resources
	Admin.AddSearchResource(product, user, order)

	// Add ActionBar
	ActionBar = action_bar.New(Admin, auth.AdminAuth{})
	ActionBar.RegisterAction(&action_bar.Action{Name: "Admin Dashboard", Link: "/admin"})

	initFuncMap()
	initRouter()
}

func sizeVariationCollection(resource interface{}, context *qor.Context) (results [][]string) {
	for _, sizeVariation := range models.SizeVariations() {
		results = append(results, []string{strconv.Itoa(int(sizeVariation.ID)), sizeVariation.Stringify()})
	}
	return
}
<|endoftext|>


<commit_before>package gannoy

import (
	"fmt"
	"math"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim

	ann := meta.treePath()

	// K := 3
	K := 50
	gannoy := GannoyIndex{
		meta:      meta,
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		g.walk(i, g.nodes.getNode(root), root, 0)
	}
}

func (g *GannoyIndex) AddItem(key int, w []float64) error {
	args := buildArgs{action: ADD, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(key int) error {
	args := buildArgs{action: DELETE, key: key, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) UpdateItem(key int, w []float64) error {
	args := buildArgs{action: UPDATE, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) GetNnsByKey(key, n, searchK int) []int {
	m := g.nodes.getNodeByKey(key)
	if !m.isLeaf() {
		return []int{}
	}
	keys := g.getAllNns(m.v, n, searchK)
	return keys
}

func (g *GannoyIndex) getAllNns(v []float64, n, searchK int) []int {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd := g.nodes.getNode(i)
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v, g.dim)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := []sorter{}
	last := -1
	for _, j := range nns {
		if j == last {
			continue
		}
		last = j
		node := g.nodes.getNode(j)
		nnsDist = append(nnsDist, sorter{value: g.distance.distance(v, node.v, g.dim), id: node.key})
	}

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	HeapSort(nnsDist, DESC, p)

	result := make([]int, p)
	for i := 0; i < p; i++ {
		result[i] = nnsDist[m-1-i].id
	}

	return result
}

func (g *GannoyIndex) addItem(key int, w []float64) error {
	n := g.nodes.newNode()
	n.key = key
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			// fmt.Printf("root: %d\n", g.meta.roots()[index])
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.nodes.maps.add(n.id, key)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	id := g.findBranchByVector(root, n.v)
	found := g.nodes.getNode(id)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, id)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{id, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		// fmt.Printf("m: %d, org_parent: %d\n", m, org_parent)
		if org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == id {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
		}
	}
}

func (g *GannoyIndex) removeItem(key int) error {
	n := g.nodes.getNodeByKey(key)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.nodes.maps.remove(key)
	n.free = true
	n.save()
	g.nodes.free.push(n.id)

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	if node.isRoot(root) {
		g.meta.updateRoot(root, -1)
		return
	}
	parent := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		if parent.isRoot(root) {
			g.meta.updateRoot(root, other)
		} else {
			grandParent := g.nodes.getNode(parent.parents[root])
			children := []int{}
			for _, child := range grandParent.children {
				if child == node.parents[root] {
					children = append(children, other)
				} else {
					children = append(children, child)
				}
			}
			grandParent.nDescendants--
			grandParent.children = children
			grandParent.save()
		}

		otherNode := g.nodes.getNode(other)
		otherNode.updateParents(root, parent.parents[root])

		parent.free = true
		parent.save()
		g.nodes.free.push(parent.id)
	}
}

func (g GannoyIndex) findBranchByVector(id int, v []float64) int {
	node := g.nodes.getNode(id)
	if node.isLeaf() || node.isBucket() {
		return id
	}
	side := g.distance.side(node, v, g.dim, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, ids []int) int {
	if len(ids) == 1 {
		n := g.nodes.getNode(ids[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return ids[0]
	}

	if len(ids) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(ids)
		m.parents[root] = parent
		m.children = ids
		m.save()
		for _, child := range ids {
			c := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(ids))
	for i, id := range ids {
		children[i] = g.nodes.getNode(id)
	}

	childrenIds := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(ids)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.dim, g.random, m)
	for _, id := range ids {
		n := g.nodes.getNode(id)
		side := g.distance.side(m, n.v, g.dim, g.random)
		childrenIds[side] = append(childrenIds[side], id)
	}

	for len(childrenIds[0]) == 0 || len(childrenIds[1]) == 0 {
		childrenIds[0] = []int{}
		childrenIds[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, id := range ids {
			side := g.random.flip()
			childrenIds[side] = append(childrenIds[side], id)
		}
	}

	var flip int
	if len(childrenIds[0]) > len(childrenIds[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIds[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	key    int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.key, args.w)
		case DELETE:
			args.result <- g.removeItem(args.key)
		case UPDATE:
			err := g.removeItem(args.key)
			if err != nil {
				args.result <- err
			} else {
				args.result <- g.addItem(args.key, args.w)
			}
		}
	}
}

func (g GannoyIndex) walk(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, node.key, node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			g.walk(root, g.nodes.getNode(child), child, tab+1)
		}
	}
}
<commit_msg>Use destroy and free.push.<commit_after>package gannoy

import (
	"fmt"
	"math"
	"sort"
	"sync"

	"github.com/gansidui/priority_queue"
)

type GannoyIndex struct {
	meta      meta
	tree      int
	dim       int
	distance  Distance
	random    Random
	nodes     Nodes
	K         int
	buildChan chan buildArgs
}

func NewGannoyIndex(metaFile string, distance Distance, random Random) (GannoyIndex, error) {

	meta, err := loadMeta(metaFile)
	if err != nil {
		return GannoyIndex{}, err
	}
	tree := meta.tree
	dim := meta.dim

	ann := meta.treePath()

	// K := 3
	K := 50
	gannoy := GannoyIndex{
		meta:      meta,
		tree:      tree,
		dim:       dim,
		distance:  distance,
		random:    random,
		K:         K,
		nodes:     newNodes(ann, tree, dim, K),
		buildChan: make(chan buildArgs, 1),
	}
	go gannoy.builder()
	return gannoy, nil
}

func (g GannoyIndex) Tree() {
	for i, root := range g.meta.roots() {
		g.walk(i, g.nodes.getNode(root), root, 0)
	}
}

func (g *GannoyIndex) AddItem(key int, w []float64) error {
	args := buildArgs{action: ADD, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) RemoveItem(key int) error {
	args := buildArgs{action: DELETE, key: key, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) UpdateItem(key int, w []float64) error {
	args := buildArgs{action: UPDATE, key: key, w: w, result: make(chan error)}
	g.buildChan <- args
	return <-args.result
}

func (g *GannoyIndex) GetNnsByKey(key, n, searchK int) []int {
	m := g.nodes.getNodeByKey(key)
	if !m.isLeaf() {
		return []int{}
	}
	keys := g.getAllNns(m.v, n, searchK)
	return keys
}

func (g *GannoyIndex) getAllNns(v []float64, n, searchK int) []int {
	if searchK == -1 {
		searchK = n * g.tree
	}

	q := priority_queue.New()
	for _, root := range g.meta.roots() {
		q.Push(&Queue{priority: math.Inf(1), value: root})
	}

	nns := []int{}
	for len(nns) < searchK && q.Len() > 0 {
		top := q.Top().(*Queue)
		d := top.priority
		i := top.value

		nd := g.nodes.getNode(i)
		q.Pop()
		if nd.isLeaf() {
			nns = append(nns, i)
		} else if nd.nDescendants <= g.K {
			dst := nd.children
			nns = append(nns, dst...)
		} else {
			margin := g.distance.margin(nd, v, g.dim)
			q.Push(&Queue{priority: math.Min(d, +margin), value: nd.children[1]})
			q.Push(&Queue{priority: math.Min(d, -margin), value: nd.children[0]})
		}
	}

	sort.Ints(nns)
	nnsDist := []sorter{}
	last := -1
	for _, j := range nns {
		if j == last {
			continue
		}
		last = j
		node := g.nodes.getNode(j)
		nnsDist = append(nnsDist, sorter{value: g.distance.distance(v, node.v, g.dim), id: node.key})
	}

	m := len(nnsDist)
	p := m
	if n < m {
		p = n
	}

	HeapSort(nnsDist, DESC, p)

	result := make([]int, p)
	for i := 0; i < p; i++ {
		result[i] = nnsDist[m-1-i].id
	}

	return result
}

func (g *GannoyIndex) addItem(key int, w []float64) error {
	n := g.nodes.newNode()
	n.key = key
	n.v = w
	n.parents = make([]int, g.tree)
	err := n.save()
	if err != nil {
		return err
	}
	// fmt.Printf("id %d\n", n.id)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for index := range buildChan {
			// fmt.Printf("root: %d\n", g.meta.roots()[index])
			g.build(index, g.meta.roots()[index], n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}

	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)
	g.nodes.maps.add(n.id, key)

	return nil
}

func (g *GannoyIndex) build(index, root int, n Node) {
	if root == -1 {
		// 最初のノード
		n.parents[index] = -1
		n.save()
		g.meta.updateRoot(index, n.id)
		return
	}
	id := g.findBranchByVector(root, n.v)
	found := g.nodes.getNode(id)
	// fmt.Printf("Found %d\n", item)

	org_parent := found.parents[index]
	if found.isBucket() && len(found.children) < g.K {
		// ノードに余裕があれば追加
		// fmt.Printf("pattern bucket\n")
		n.updateParents(index, id)
		found.nDescendants++
		found.children = append(found.children, n.id)
		found.save()
	} else {
		// ノードが上限またはリーフノードであれば新しいノードを追加
		willDelete := false
		var indices []int
		if found.isLeaf() {
			// fmt.Printf("pattern leaf node\n")
			indices = []int{id, n.id}
		} else {
			// fmt.Printf("pattern full backet\n")
			indices = append(found.children, n.id)
			willDelete = true
		}

		m := g.makeTree(index, org_parent, indices)
		// fmt.Printf("m: %d, org_parent: %d\n", m, org_parent)
		if org_parent == -1 {
			// rootノードの入れ替え
			g.meta.updateRoot(index, m)
		} else {
			parent := g.nodes.getNode(org_parent)
			parent.nDescendants++
			children := make([]int, len(parent.children))
			for i, child := range parent.children {
				if child == id {
					// 新しいノードに変更
					children[i] = m
				} else {
					// 既存のノードのまま
					children[i] = child
				}
			}
			parent.children = children
			parent.save()

		}
		if willDelete {
			found.destroy()
			g.nodes.free.push(found.id)
		}
	}
}

func (g *GannoyIndex) removeItem(key int) error {
	n := g.nodes.getNodeByKey(key)

	var wg sync.WaitGroup
	wg.Add(g.tree)
	buildChan := make(chan int, g.tree)
	worker := func(n Node) {
		for root := range buildChan {
			g.remove(root, n)
			wg.Done()
		}
	}

	for i := 0; i < 3; i++ {
		go worker(n)
	}
	for index, _ := range g.meta.roots() {
		buildChan <- index
	}

	wg.Wait()
	close(buildChan)

	g.nodes.maps.remove(key)
	n.destroy()
	g.nodes.free.push(n.id)

	return nil
}

func (g *GannoyIndex) remove(root int, node Node) {
	if node.isRoot(root) {
		g.meta.updateRoot(root, -1)
		return
	}
	parent := g.nodes.getNode(node.parents[root])
	if parent.isBucket() && len(parent.children) > 2 {
		// fmt.Printf("pattern bucket\n")
		target := -1
		for i, child := range parent.children {
			if child == node.id {
				target = i
			}
		}
		if target == -1 {
			return
		}
		children := append(parent.children[:target], parent.children[(target+1):]...)
		parent.nDescendants--
		parent.children = children
		parent.save()
	} else {
		// fmt.Printf("pattern leaf node\n")
		var other int
		for _, child := range parent.children {
			if child != node.id {
				other = child
			}
		}
		if parent.isRoot(root) {
			g.meta.updateRoot(root, other)
		} else {
			grandParent := g.nodes.getNode(parent.parents[root])
			children := []int{}
			for _, child := range grandParent.children {
				if child == node.parents[root] {
					children = append(children, other)
				} else {
					children = append(children, child)
				}
			}
			grandParent.nDescendants--
			grandParent.children = children
			grandParent.save()
		}

		otherNode := g.nodes.getNode(other)
		otherNode.updateParents(root, parent.parents[root])

		parent.destroy()
		g.nodes.free.push(parent.id)
	}
}

func (g GannoyIndex) findBranchByVector(id int, v []float64) int {
	node := g.nodes.getNode(id)
	if node.isLeaf() || node.isBucket() {
		return id
	}
	side := g.distance.side(node, v, g.dim, g.random)
	return g.findBranchByVector(node.children[side], v)
}

func (g *GannoyIndex) makeTree(root, parent int, ids []int) int {
	if len(ids) == 1 {
		n := g.nodes.getNode(ids[0])
		if len(n.parents) == 0 {
			n.parents = make([]int, g.tree)
		}
		n.updateParents(root, parent)
		return ids[0]
	}

	if len(ids) <= g.K {
		m := g.nodes.newNode()
		m.parents = make([]int, g.tree)
		m.nDescendants = len(ids)
		m.parents[root] = parent
		m.children = ids
		m.save()
		for _, child := range ids {
			c := g.nodes.getNode(child)
			if len(c.parents) == 0 {
				c.parents = make([]int, g.tree)
			}
			c.updateParents(root, m.id)
		}
		return m.id
	}

	children := make([]Node, len(ids))
	for i, id := range ids {
		children[i] = g.nodes.getNode(id)
	}

	childrenIds := [2][]int{[]int{}, []int{}}

	m := g.nodes.newNode()
	m.parents = make([]int, g.tree)
	m.nDescendants = len(ids)
	m.parents[root] = parent

	m = g.distance.createSplit(children, g.dim, g.random, m)
	for _, id := range ids {
		n := g.nodes.getNode(id)
		side := g.distance.side(m, n.v, g.dim, g.random)
		childrenIds[side] = append(childrenIds[side], id)
	}

	for len(childrenIds[0]) == 0 || len(childrenIds[1]) == 0 {
		childrenIds[0] = []int{}
		childrenIds[1] = []int{}
		for z := 0; z < g.dim; z++ {
			m.v[z] = 0.0
		}
		for _, id := range ids {
			side := g.random.flip()
			childrenIds[side] = append(childrenIds[side], id)
		}
	}

	var flip int
	if len(childrenIds[0]) > len(childrenIds[1]) {
		flip = 1
	}

	m.save()
	for side := 0; side < 2; side++ {
		m.children[side^flip] = g.makeTree(root, m.id, childrenIds[side^flip])
	}
	m.save()

	return m.id
}

type buildArgs struct {
	action int
	key    int
	w      []float64
	result chan error
}

func (g *GannoyIndex) builder() {
	for args := range g.buildChan {
		switch args.action {
		case ADD:
			args.result <- g.addItem(args.key, args.w)
		case DELETE:
			args.result <- g.removeItem(args.key)
		case UPDATE:
			err := g.removeItem(args.key)
			if err != nil {
				args.result <- err
			} else {
				args.result <- g.addItem(args.key, args.w)
			}
		}
	}
}

func (g GannoyIndex) walk(root int, node Node, id, tab int) {
	for i := 0; i < tab*2; i++ {
		fmt.Print(" ")
	}
	fmt.Printf("%d [%d] (%d) [nDescendants: %d, v: %v]\n", id, node.key, node.parents[root], node.nDescendants, node.v)
	if !node.isLeaf() {
		for _, child := range node.children {
			g.walk(root, g.nodes.getNode(child), child, tab+1)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"strconv"
	"sync"
	"time"

	"github.com/giagiannis/data-profiler/core"
)

// TaskEngine is deployed once for the server's lifetime and keeps the tasks
// which are executed.
type TaskEngine struct {
	Tasks []*Task
	lock  sync.Mutex
}

// NewTaskEngine initializes a new TasEngine object. Called once per server deployment.
func NewTaskEngine() *TaskEngine {
	te := new(TaskEngine)
	te.lock = *new(sync.Mutex)
	return te
}

// Submit appends a new task to the task engine and initializes its execution.
func (e *TaskEngine) Submit(t *Task) {
	if t == nil {
		return
	}
	e.lock.Lock()
	go t.Run()
	e.Tasks = append(e.Tasks, t)
	e.lock.Unlock()
}

// Task is the primitive struct that represents a task of the server.
type Task struct {
	Status      string
	Started     time.Time
	Duration    float64
	Description string
	Dataset     *ModelDataset
	fnc         func() error
}

// Run is responsible to execute to task's method and update the task status
// accordingly.
func (t *Task) Run() {
	t.Started = time.Now()
	t.Status = "RUNNING"
	err := t.fnc()
	if err != nil {
		t.Status = "ERROR - " + err.Error()
	} else {
		t.Status = "DONE"
		t.Duration = time.Since(t.Started).Seconds()
	}
}

// NewSMComputationTask initializes a new Similarity Matrix computation task.
func NewSMComputationTask(datasetID string, conf map[string]string) *Task {
	task := new(Task)
	dts := modelDatasetGetInfo(datasetID)
	task.Dataset = dts
	task.Description = fmt.Sprintf("SM Computation for %s, type %s\n",
		dts.Name, conf["estimatorType"])
	task.fnc = func() error {
		datasets := core.DiscoverDatasets(dts.Path)
		estType := *core.NewDatasetSimilarityEstimatorType(conf["estimatorType"])
		est := core.NewDatasetSimilarityEstimator(estType, datasets)
		est.Configure(conf)
		if conf["popPolicy"] == "aprx" {
			pop := new(core.DatasetSimilarityPopulationPolicy)
			pop.PolicyType = core.PopulationPolicyAprx
			val, err := strconv.ParseFloat(conf["popParameterValue"], 64)
			if err != nil {
				log.Println(err)
			}
			if conf["popParameter"] == "popCount" {
				pop.Parameters = map[string]float64{"count": val}
			}
			if conf["popParameter"] == "popThreshold" {
				pop.Parameters = map[string]float64{"threshold": val}
			}
			est.SetPopulationPolicy(*pop)
		}
		err := est.Compute()
		if err != nil {
			return err
		}
		sm := est.SimilarityMatrix()
		//		var smID string
		if sm != nil {
			//smID =
			modelSimilarityMatrixInsert(datasetID, sm.Serialize(), est.Serialize(), conf)
		}
		//modelEstimatorInsert(datasetID, smID, est.Serialize(), conf)
		return nil
	}
	return task
}

// NewMDSComputationTask initializes a new Multidimensional Scaling execution task.
func NewMDSComputationTask(smID, datasetID string, conf map[string]string) *Task {
	smModel := modelSimilarityMatrixGet(smID)
	if smModel == nil {
		log.Println("SM not found")
		return nil
	}

	cnt, err := ioutil.ReadFile(smModel.Path)
	if err != nil {
		log.Println(err)
	}
	sm := new(core.DatasetSimilarityMatrix)
	sm.Deserialize(cnt)
	k, err := strconv.ParseInt(conf["k"], 10, 64)
	if err != nil {
		log.Println(err)
	}

	dat := modelDatasetGetInfo(datasetID)
	task := new(Task)
	task.Dataset = dat
	task.Description = fmt.Sprintf("MDS Execution for %s with k=%d\n",
		dat.Name, k)
	task.fnc = func() error {
		mds := core.NewMDScaling(sm, int(k), Conf.Scripts.MDS)
		err = mds.Compute()
		if err != nil {
			return err
		}
		gof := fmt.Sprintf("%.5f", mds.Gof())
		stress := fmt.Sprintf("%.5f", mds.Stress())
		modelCoordinatesInsert(mds.Coordinates(), dat.ID, conf["k"], gof, stress, smID)
		return nil
	}
	return task
}

// NewOperatorRunTask initializes a new operator execution task.
func NewOperatorRunTask(operatorID string) *Task {
	m := modelOperatorGet(operatorID)
	if m == nil {
		log.Println("Operator was not found")
		return nil
	}
	dat := modelDatasetGetInfo(m.DatasetID)
	for _, f := range modelDatasetGetFiles(dat.ID) {
		dat.Files = append(dat.Files, dat.Path+"/"+f)
	}
	task := new(Task)
	task.Description = fmt.Sprintf("%s evaluation", m.Name)
	task.Dataset = dat
	task.fnc = func() error {
		eval, err := core.NewDatasetEvaluator(core.OnlineEval,
			map[string]string{
				"script":  m.Path,
				"testset": "",
			})
		if err != nil {
			log.Println(err)
		}
		scores := core.NewDatasetScores()
		for _, f := range dat.Files {
			s, err := eval.Evaluate(f)
			if err != nil {
				log.Println(err)
			} else {
				scores.Scores[path.Base(f)] = s
			}
		}
		cnt, _ := scores.Serialize()
		modelOperatorScoresInsert(operatorID, cnt)
		return nil
	}
	return task
}

// NewModelTrainTask generates a new ML for a given (dataset,operator) combination,
// according to the user-specified parameters.
func NewModelTrainTask(datasetID, operatorID string, sr float64,
	modelType string,
	coordinatesID, mlScript string,
	matrixID, k string) *Task {
	m := modelDatasetGetInfo(datasetID)
	task := new(Task)
	task.Description = fmt.Sprintf("Model training (%s for %s)", path.Base(mlScript), m.Name)
	task.Dataset = m
	task.fnc = func() error {
		datasets := core.DiscoverDatasets(m.Path)
		o := modelOperatorGet(operatorID)
		var evaluator core.DatasetEvaluator
		var err error
		if o.ScoresFile != "" {
			evaluator, err = core.NewDatasetEvaluator(core.FileBasedEval, map[string]string{"scores": o.ScoresFile})
		} else {
			evaluator, err = core.NewDatasetEvaluator(core.OnlineEval, map[string]string{"script": o.Path, "testset": ""})
		}
		if err != nil {
			log.Println(err)
			return err
		}
		t := core.NewModelerType(modelType)
		modeler := core.NewModeler(t, datasets, sr, evaluator)
		var conf map[string]string
		if t == core.ScriptBasedModelerType {
			c := modelCoordinatesGet(coordinatesID)
			conf = map[string]string{"script": mlScript, "coordinates": path.Base(c.Path)}
		} else if t == core.KNNModelerType {
			m := modelSimilarityMatrixGet(matrixID)
			conf = map[string]string{"k": k, "smatrix": path.Base(m.Path)}
		}
		modeler.Configure(conf)
		err = modeler.Run()
		if err != nil {
			log.Println(err)
			return err
		}
		// serialze appxValues
		var cnt [][]float64
		cnt = append(cnt, modeler.AppxValues())
		appxBuffer := serializeCSVFile(cnt)
		samplesBuffer, _ := json.Marshal(modeler.Samples())
		errors := make(map[string]string)
		for k, v := range modeler.ErrorMetrics() {
			errors[k] = fmt.Sprintf("%.5f", v)
		}
		modelDatasetModelInsert(coordinatesID, operatorID, datasetID, samplesBuffer, appxBuffer, conf, errors, sr)
		return nil
	}
	return task
}
<commit_msg>reverting<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"strconv"
	"sync"
	"time"

	"github.com/giagiannis/data-profiler/core"
)

// TaskEngine is deployed once for the server's lifetime and keeps the tasks
// which are executed.
type TaskEngine struct {
	Tasks []*Task
	lock  sync.Mutex
}

// NewTaskEngine initializes a new TasEngine object. Called once per server deployment.
func NewTaskEngine() *TaskEngine {
	te := new(TaskEngine)
	te.lock = *new(sync.Mutex)
	return te
}

// Submit appends a new task to the task engine and initializes its execution.
func (e *TaskEngine) Submit(t *Task) {
	if t == nil {
		return
	}
	e.lock.Lock()
	go t.Run()
	e.Tasks = append(e.Tasks, t)
	e.lock.Unlock()
}

// Task is the primitive struct that represents a task of the server.
type Task struct {
	Status      string
	Started     time.Time
	Duration    float64
	Description string
	Dataset     *ModelDataset
	fnc         func() error
}

// Run is responsible to execute to task's method and update the task status
// accordingly.
func (t *Task) Run() {
	t.Started = time.Now()
	t.Status = "RUNNING"
	err := t.fnc()
	if err != nil {
		t.Status = "ERROR - " + err.Error()
	} else {
		t.Status = "DONE"
		t.Duration = time.Since(t.Started).Seconds()
	}
}

// NewSMComputationTask initializes a new Similarity Matrix computation task.
func NewSMComputationTask(datasetID string, conf map[string]string) *Task {
	task := new(Task)
	dts := modelDatasetGetInfo(datasetID)
	task.Dataset = dts
	task.Description = fmt.Sprintf("SM Computation for %s, type %s\n",
		dts.Name, conf["estimatorType"])
	task.fnc = func() error {
		datasets := core.DiscoverDatasets(dts.Path)
		estType := *core.NewDatasetSimilarityEstimatorType(conf["estimatorType"])
		est := core.NewDatasetSimilarityEstimator(estType, datasets)
		est.Configure(conf)
		if conf["popPolicy"] == "aprx" {
			pop := new(core.DatasetSimilarityPopulationPolicy)
			pop.PolicyType = core.PopulationPolicyAprx
			val, err := strconv.ParseFloat(conf["popParameterValue"], 64)
			if err != nil {
				log.Println(err)
			}
			if conf["popParameter"] == "popCount" {
				pop.Parameters = map[string]float64{"count": val}
			}
			if conf["popParameter"] == "popThreshold" {
				pop.Parameters = map[string]float64{"threshold": val}
			}
			est.SetPopulationPolicy(*pop)
		}
		err := est.Compute()
		if err != nil {
			return err
		}
		sm := est.SimilarityMatrix()
		//		var smID string
		if sm != nil {
			//smID =
			modelSimilarityMatrixInsert(datasetID, sm.Serialize(), est.Serialize(), conf)
		}
		//modelEstimatorInsert(datasetID, smID, est.Serialize(), conf)
		return nil
	}
	return task
}

// NewMDSComputationTask initializes a new Multidimensional Scaling execution task.
func NewMDSComputationTask(smID, datasetID string, conf map[string]string) *Task {
	smModel := modelSimilarityMatrixGet(smID)
	if smModel == nil {
		log.Println("SM not found")
		return nil
	}

	cnt, err := ioutil.ReadFile(smModel.Path)
	if err != nil {
		log.Println(err)
	}
	sm := new(core.DatasetSimilarityMatrix)
	sm.Deserialize(cnt)
	k, err := strconv.ParseInt(conf["k"], 10, 64)
	if err != nil {
		log.Println(err)
	}

	dat := modelDatasetGetInfo(datasetID)
	task := new(Task)
	task.Dataset = dat
	task.Description = fmt.Sprintf("MDS Execution for %s with k=%d\n",
		dat.Name, k)
	task.fnc = func() error {
		mds := core.NewMDScaling(sm, int(k), Conf.Scripts.MDS)
		err = mds.Compute()
		if err != nil {
			return err
		}
		gof := fmt.Sprintf("%.5f", mds.Gof())
		stress := fmt.Sprintf("%.5f", mds.Stress())
		modelCoordinatesInsert(mds.Coordinates(), dat.ID, conf["k"], gof, stress, smID)
		return nil
	}
	return task
}

// NewOperatorRunTask initializes a new operator execution task.
func NewOperatorRunTask(operatorID string) *Task {
	m := modelOperatorGet(operatorID)
	if m == nil {
		log.Println("Operator was not found")
		return nil
	}
	dat := modelDatasetGetInfo(m.DatasetID)
	for _, f := range modelDatasetGetFiles(dat.ID) {
		dat.Files = append(dat.Files, dat.Path+"/"+f)
	}
	task := new(Task)
	task.Description = fmt.Sprintf("%s evaluation", m.Name)
	task.Dataset = dat
	task.fnc = func() error {
		eval, err := core.NewDatasetEvaluator(core.OnlineEval,
			map[string]string{
				"script":  m.Path,
				"testset": "",
			})
		if err != nil {
			log.Println(err)
		}
		scores := core.NewDatasetScores()
		for _, f := range dat.Files {
			s, err := eval.Evaluate(f)
			if err != nil {
				log.Println(err)
			} else {
				scores.Scores[path.Base(f)] = s
			}
		}
		cnt, _ := scores.Serialize()
		modelOperatorScoresInsert(operatorID, cnt)
		return nil
	}
	return task
}

// NewModelTrainTask generates a new ML for a given (dataset,operator) combination,
// according to the user-specified parameters.
func NewModelTrainTask(datasetID, operatorID string, sr float64,
	modelType string,
	coordinatesID, mlScript string,
	matrixID, k string) *Task {
	m := modelDatasetGetInfo(datasetID)
	task := new(Task)
	task.Description = fmt.Sprintf("Model training (%s for %s)", path.Base(mlScript), m.Name)
	task.Dataset = m
	task.fnc = func() error {
		datasets := core.DiscoverDatasets(m.Path)
		o := modelOperatorGet(operatorID)
		var evaluator core.DatasetEvaluator
		var err error
		if o.ScoresFile != "" {
			evaluator, err = core.NewDatasetEvaluator(core.FileBasedEval, map[string]string{"scores": o.ScoresFile})
		} else {
			evaluator, err = core.NewDatasetEvaluator(core.OnlineEval, map[string]string{"script": o.Path, "testset": ""})
		}
		if err != nil {
			log.Println(err)
			return err
		}
		t := core.NewModelerType(modelType)
		modeler := core.NewModeler(t, datasets, sr, evaluator)
		var conf map[string]string
		if t == core.ScriptBasedModelerType {
			c := modelCoordinatesGet(coordinatesID)
			conf = map[string]string{"script": mlScript, "coordinates": c.Path}
		} else if t == core.KNNModelerType {
			m := modelSimilarityMatrixGet(matrixID)
			conf = map[string]string{"k": k, "smatrix": m.Path}
		}
		modeler.Configure(conf)
		err = modeler.Run()
		if err != nil {
			log.Println(err)
			return err
		}
		// serialze appxValues
		var cnt [][]float64
		cnt = append(cnt, modeler.AppxValues())
		appxBuffer := serializeCSVFile(cnt)
		samplesBuffer, _ := json.Marshal(modeler.Samples())
		errors := make(map[string]string)
		for k, v := range modeler.ErrorMetrics() {
			errors[k] = fmt.Sprintf("%.5f", v)
		}
		modelDatasetModelInsert(coordinatesID, operatorID, datasetID, samplesBuffer, appxBuffer, conf, errors, sr)
		return nil
	}
	return task
}
<|endoftext|>


<commit_before>// Copyright 2017 tsuru 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 integration

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strings"
	"time"

	"gopkg.in/check.v1"
)

var (
	T            = NewCommand("tsuru").WithArgs
	allPlatforms = []string{
		"tsuru/python",
		"tsuru/go",
		"tsuru/buildpack",
		"tsuru/cordova",
		"tsuru/elixir",
		"tsuru/java",
		"tsuru/nodejs",
		"tsuru/php",
		"tsuru/play",
		"tsuru/pypy",
		"tsuru/python3",
		"tsuru/ruby",
		"tsuru/static",
	}
	allProvisioners = []string{
		"docker",
		"swarm",
	}
	flows = []ExecFlow{
		platformsToInstall(),
		installerConfigTest(),
		installerTest(),
		targetTest(),
		loginTest(),
		removeInstallNodes(),
		quotaTest(),
		teamTest(),
		poolAdd(),
		platformAdd(),
		exampleApps(),
	}
)

var installerConfig = `driver:
  name: virtualbox
  options:
    virtualbox-cpu-count: 2
    virtualbox-memory: 2048
hosts:
  apps:
    size: 2
components:
  tsuru:
    version: latest
    install-dashboard: false
`

func platformsToInstall() ExecFlow {
	flow := ExecFlow{
		provides: []string{"platformimages"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		for _, platImg := range allPlatforms {
			env.Add("platformimages", platImg)
		}
	}
	return flow
}

func installerConfigTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"installerconfig"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		f, err := ioutil.TempFile("", "installer-config")
		c.Assert(err, check.IsNil)
		defer f.Close()
		f.Write([]byte(installerConfig))
		env.Set("installerconfig", f.Name())
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := NewCommand("rm", "{{.installerconfig}}").Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func installerTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"targetaddr"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("install", "--config", "{{.installerconfig}}").WithTimeout(30 * time.Minute).Run(env)
		c.Assert(res, ResultOk)
		regex := regexp.MustCompile(`(?si).*Core Hosts:.*?([\d.]+)\s.*`)
		parts := regex.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		targetHost := parts[1]
		regex = regexp.MustCompile(`(?si).*Tsuru API.*?\|\s(\d+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		targetPort := parts[1]
		env.Set("targetaddr", fmt.Sprintf("http://%s:%s", targetHost, targetPort))
		regex = regexp.MustCompile(`\| (https?[^\s]+?) \|`)
		allParts := regex.FindAllStringSubmatch(res.Stdout.String(), -1)
		for _, parts = range allParts {
			c.Assert(parts, check.HasLen, 2)
			env.Add("nodeopts", fmt.Sprintf("--register address=%s --cacert ~/.tsuru/installs/tsuru/certs/ca.pem --clientcert ~/.tsuru/installs/tsuru/certs/cert.pem --clientkey ~/.tsuru/installs/tsuru/certs/key.pem", parts[1]))
			env.Add("nodestoremove", parts[1])
		}
		regex = regexp.MustCompile(`Username: ([[:print:]]+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		env.Set("adminuser", parts[1])
		regex = regexp.MustCompile(`Password: ([[:print:]]+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		env.Set("adminpassword", parts[1])
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("uninstall", "-y").Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func targetTest() ExecFlow {
	flow := ExecFlow{}
	flow.forward = func(c *check.C, env *Environment) {
		targetName := "integration-target"
		res := T("target-add", targetName, "{{.targetaddr}}").Run(env)
		c.Assert(res, ResultOk)
		res = T("target-list").Run(env)
		c.Assert(res, ResultMatches, Expected{Stdout: `\s+` + targetName + ` .*`})
		res = T("target-set", targetName).Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func loginTest() ExecFlow {
	flow := ExecFlow{}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("login", "{{.adminuser}}").WithInput("{{.adminpassword}}").Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func removeInstallNodes() ExecFlow {
	flow := ExecFlow{
		matrix: map[string]string{
			"node": "nodestoremove",
		},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("node-remove", "-y", "--no-rebalance", "{{.node}}").Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func quotaTest() ExecFlow {
	flow := ExecFlow{
		requires: []string{"adminuser"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("user-quota-change", "{{.adminuser}}", "100").Run(env)
		c.Assert(res, ResultOk)
		res = T("user-quota-view", "{{.adminuser}}").Run(env)
		c.Assert(res, ResultMatches, Expected{Stdout: `(?s)Apps usage.*/100`})
	}
	return flow
}

func teamTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"team"},
	}
	teamName := "integration-team"
	flow.forward = func(c *check.C, env *Environment) {
		res := T("team-create", teamName).Run(env)
		c.Assert(res, ResultOk)
		env.Set("team", teamName)
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("team-remove", "-y", teamName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func poolAdd() ExecFlow {
	flow := ExecFlow{
		provides: []string{"poolnames"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		for _, prov := range allProvisioners {
			poolName := "ipool-" + prov
			res := T("pool-add", "--provisioner", prov, poolName).Run(env)
			c.Assert(res, ResultOk)
			env.Add("poolnames", poolName)
			res = T("pool-teams-add", poolName, "{{.team}}").Run(env)
			c.Assert(res, ResultOk)
			res = T("node-add", "{{.nodeopts}}", "pool="+poolName).Run(env)
			c.Assert(res, ResultOk)
			res = T("event-list").Run(env)
			c.Assert(res, ResultOk)
			nodeopts := env.All("nodeopts")
			env.Set("nodeopts", append(nodeopts[1:], nodeopts[0])...)
			regex := regexp.MustCompile(`node.create.*?node:\s+(.*?)\s+`)
			parts := regex.FindStringSubmatch(res.Stdout.String())
			c.Assert(parts, check.HasLen, 2)
			env.Add("nodeaddrs", parts[1])
			regex = regexp.MustCompile(parts[1] + `.*?ready`)
			ok := retry(time.Minute, func() bool {
				res = T("node-list").Run(env)
				return regex.MatchString(res.Stdout.String())
			})
			c.Assert(ok, check.Equals, true, check.Commentf("node not ready after 1 minute: %v", res))
		}
	}
	flow.backward = func(c *check.C, env *Environment) {
		for _, node := range env.All("nodeaddrs") {
			res := T("node-remove", "-y", "--no-rebalance", node).Run(env)
			c.Check(res, ResultOk)
		}
		for _, prov := range allProvisioners {
			poolName := "ipool-" + prov
			res := T("pool-teams-remove", poolName, "{{.team}}").Run(env)
			c.Check(res, ResultOk)
			res = T("pool-remove", "-y", poolName).Run(env)
			c.Check(res, ResultOk)
		}
	}
	return flow
}

func platformAdd() ExecFlow {
	flow := ExecFlow{
		provides: []string{"platforms"},
		matrix: map[string]string{
			"platimg": "platformimages",
		},
		parallel: true,
	}
	flow.forward = func(c *check.C, env *Environment) {
		img := env.Get("platimg")
		suffix := img[strings.LastIndex(img, "/")+1:]
		platName := "iplat-" + suffix
		res := T("platform-add", platName, "-i", img).WithTimeout(15 * time.Minute).Run(env)
		c.Assert(res, ResultOk)
		env.Add("platforms", platName)
		res = T("platform-list").Run(env)
		c.Assert(res, ResultOk)
		c.Assert(res, ResultMatches, Expected{Stdout: "(?s).*- " + platName + ".*"})
	}
	flow.backward = func(c *check.C, env *Environment) {
		img := env.Get("platimg")
		suffix := img[strings.LastIndex(img, "/")+1:]
		platName := "iplat-" + suffix
		res := T("platform-remove", "-y", platName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func exampleApps() ExecFlow {
	flow := ExecFlow{
		matrix: map[string]string{
			"pool": "poolnames",
			"plat": "platforms",
		},
		parallel: true,
	}
	appName := "iapp-{{.plat}}-{{.pool}}"
	flow.forward = func(c *check.C, env *Environment) {
		res := T("app-create", appName, "{{.plat}}", "-t", "{{.team}}", "-o", "{{.pool}}").Run(env)
		c.Assert(res, ResultOk)
		res = T("app-info", "-a", appName).Run(env)
		c.Assert(res, ResultOk)
		platRE := regexp.MustCompile(`(?s)Platform: (.*?)\n`)
		parts := platRE.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		lang := strings.Replace(parts[1], "iplat-", "", -1)
		res = T("app-deploy", "-a", appName, "{{.examplesdir}}/"+lang+"/").Run(env)
		c.Assert(res, ResultOk)
		res = T("app-info", "-a", appName).Run(env)
		c.Assert(res, ResultOk)
		addrRE := regexp.MustCompile(`(?s)Address: (.*?)\n`)
		parts = addrRE.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		cmd := NewCommand("curl", "-sSf", "http://"+parts[1])
		ok := retry(time.Minute, func() bool {
			res = cmd.Run(env)
			return res.ExitCode == 0
		})
		c.Assert(ok, check.Equals, true, check.Commentf("invalid result: %v", res))
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("app-remove", "-y", "-a", appName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func (s *S) TestBase(c *check.C) {
	env := NewEnvironment()
	if !env.Has("enabled") {
		return
	}
	var executedFlows []*ExecFlow
	defer func() {
		for i := len(executedFlows) - 1; i >= 0; i-- {
			executedFlows[i].Rollback(c, env)
		}
	}()
	for i := range flows {
		f := &flows[i]
		if len(f.provides) > 0 {
			providesAll := true
			for _, envVar := range f.provides {
				if env.Get(envVar) == "" {
					providesAll = false
					break
				}
			}
			if providesAll {
				continue
			}
		}
		executedFlows = append(executedFlows, f)
		f.Run(c, env)
	}
}
<commit_msg>integration: increase time waiting for application to be available<commit_after>// Copyright 2017 tsuru 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 integration

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strings"
	"time"

	"gopkg.in/check.v1"
)

var (
	T            = NewCommand("tsuru").WithArgs
	allPlatforms = []string{
		"tsuru/python",
		"tsuru/go",
		"tsuru/buildpack",
		"tsuru/cordova",
		"tsuru/elixir",
		"tsuru/java",
		"tsuru/nodejs",
		"tsuru/php",
		"tsuru/play",
		"tsuru/pypy",
		"tsuru/python3",
		"tsuru/ruby",
		"tsuru/static",
	}
	allProvisioners = []string{
		"docker",
		"swarm",
	}
	flows = []ExecFlow{
		platformsToInstall(),
		installerConfigTest(),
		installerTest(),
		targetTest(),
		loginTest(),
		removeInstallNodes(),
		quotaTest(),
		teamTest(),
		poolAdd(),
		platformAdd(),
		exampleApps(),
	}
)

var installerConfig = `driver:
  name: virtualbox
  options:
    virtualbox-cpu-count: 2
    virtualbox-memory: 2048
hosts:
  apps:
    size: 2
components:
  tsuru:
    version: latest
    install-dashboard: false
`

func platformsToInstall() ExecFlow {
	flow := ExecFlow{
		provides: []string{"platformimages"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		for _, platImg := range allPlatforms {
			env.Add("platformimages", platImg)
		}
	}
	return flow
}

func installerConfigTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"installerconfig"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		f, err := ioutil.TempFile("", "installer-config")
		c.Assert(err, check.IsNil)
		defer f.Close()
		f.Write([]byte(installerConfig))
		env.Set("installerconfig", f.Name())
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := NewCommand("rm", "{{.installerconfig}}").Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func installerTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"targetaddr"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("install", "--config", "{{.installerconfig}}").WithTimeout(30 * time.Minute).Run(env)
		c.Assert(res, ResultOk)
		regex := regexp.MustCompile(`(?si).*Core Hosts:.*?([\d.]+)\s.*`)
		parts := regex.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		targetHost := parts[1]
		regex = regexp.MustCompile(`(?si).*Tsuru API.*?\|\s(\d+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		targetPort := parts[1]
		env.Set("targetaddr", fmt.Sprintf("http://%s:%s", targetHost, targetPort))
		regex = regexp.MustCompile(`\| (https?[^\s]+?) \|`)
		allParts := regex.FindAllStringSubmatch(res.Stdout.String(), -1)
		for _, parts = range allParts {
			c.Assert(parts, check.HasLen, 2)
			env.Add("nodeopts", fmt.Sprintf("--register address=%s --cacert ~/.tsuru/installs/tsuru/certs/ca.pem --clientcert ~/.tsuru/installs/tsuru/certs/cert.pem --clientkey ~/.tsuru/installs/tsuru/certs/key.pem", parts[1]))
			env.Add("nodestoremove", parts[1])
		}
		regex = regexp.MustCompile(`Username: ([[:print:]]+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		env.Set("adminuser", parts[1])
		regex = regexp.MustCompile(`Password: ([[:print:]]+)`)
		parts = regex.FindStringSubmatch(res.Stdout.String())
		env.Set("adminpassword", parts[1])
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("uninstall", "-y").Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func targetTest() ExecFlow {
	flow := ExecFlow{}
	flow.forward = func(c *check.C, env *Environment) {
		targetName := "integration-target"
		res := T("target-add", targetName, "{{.targetaddr}}").Run(env)
		c.Assert(res, ResultOk)
		res = T("target-list").Run(env)
		c.Assert(res, ResultMatches, Expected{Stdout: `\s+` + targetName + ` .*`})
		res = T("target-set", targetName).Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func loginTest() ExecFlow {
	flow := ExecFlow{}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("login", "{{.adminuser}}").WithInput("{{.adminpassword}}").Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func removeInstallNodes() ExecFlow {
	flow := ExecFlow{
		matrix: map[string]string{
			"node": "nodestoremove",
		},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("node-remove", "-y", "--no-rebalance", "{{.node}}").Run(env)
		c.Assert(res, ResultOk)
	}
	return flow
}

func quotaTest() ExecFlow {
	flow := ExecFlow{
		requires: []string{"adminuser"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		res := T("user-quota-change", "{{.adminuser}}", "100").Run(env)
		c.Assert(res, ResultOk)
		res = T("user-quota-view", "{{.adminuser}}").Run(env)
		c.Assert(res, ResultMatches, Expected{Stdout: `(?s)Apps usage.*/100`})
	}
	return flow
}

func teamTest() ExecFlow {
	flow := ExecFlow{
		provides: []string{"team"},
	}
	teamName := "integration-team"
	flow.forward = func(c *check.C, env *Environment) {
		res := T("team-create", teamName).Run(env)
		c.Assert(res, ResultOk)
		env.Set("team", teamName)
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("team-remove", "-y", teamName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func poolAdd() ExecFlow {
	flow := ExecFlow{
		provides: []string{"poolnames"},
	}
	flow.forward = func(c *check.C, env *Environment) {
		for _, prov := range allProvisioners {
			poolName := "ipool-" + prov
			res := T("pool-add", "--provisioner", prov, poolName).Run(env)
			c.Assert(res, ResultOk)
			env.Add("poolnames", poolName)
			res = T("pool-teams-add", poolName, "{{.team}}").Run(env)
			c.Assert(res, ResultOk)
			res = T("node-add", "{{.nodeopts}}", "pool="+poolName).Run(env)
			c.Assert(res, ResultOk)
			res = T("event-list").Run(env)
			c.Assert(res, ResultOk)
			nodeopts := env.All("nodeopts")
			env.Set("nodeopts", append(nodeopts[1:], nodeopts[0])...)
			regex := regexp.MustCompile(`node.create.*?node:\s+(.*?)\s+`)
			parts := regex.FindStringSubmatch(res.Stdout.String())
			c.Assert(parts, check.HasLen, 2)
			env.Add("nodeaddrs", parts[1])
			regex = regexp.MustCompile(parts[1] + `.*?ready`)
			ok := retry(time.Minute, func() bool {
				res = T("node-list").Run(env)
				return regex.MatchString(res.Stdout.String())
			})
			c.Assert(ok, check.Equals, true, check.Commentf("node not ready after 1 minute: %v", res))
		}
	}
	flow.backward = func(c *check.C, env *Environment) {
		for _, node := range env.All("nodeaddrs") {
			res := T("node-remove", "-y", "--no-rebalance", node).Run(env)
			c.Check(res, ResultOk)
		}
		for _, prov := range allProvisioners {
			poolName := "ipool-" + prov
			res := T("pool-teams-remove", poolName, "{{.team}}").Run(env)
			c.Check(res, ResultOk)
			res = T("pool-remove", "-y", poolName).Run(env)
			c.Check(res, ResultOk)
		}
	}
	return flow
}

func platformAdd() ExecFlow {
	flow := ExecFlow{
		provides: []string{"platforms"},
		matrix: map[string]string{
			"platimg": "platformimages",
		},
		parallel: true,
	}
	flow.forward = func(c *check.C, env *Environment) {
		img := env.Get("platimg")
		suffix := img[strings.LastIndex(img, "/")+1:]
		platName := "iplat-" + suffix
		res := T("platform-add", platName, "-i", img).WithTimeout(15 * time.Minute).Run(env)
		c.Assert(res, ResultOk)
		env.Add("platforms", platName)
		res = T("platform-list").Run(env)
		c.Assert(res, ResultOk)
		c.Assert(res, ResultMatches, Expected{Stdout: "(?s).*- " + platName + ".*"})
	}
	flow.backward = func(c *check.C, env *Environment) {
		img := env.Get("platimg")
		suffix := img[strings.LastIndex(img, "/")+1:]
		platName := "iplat-" + suffix
		res := T("platform-remove", "-y", platName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func exampleApps() ExecFlow {
	flow := ExecFlow{
		matrix: map[string]string{
			"pool": "poolnames",
			"plat": "platforms",
		},
		parallel: true,
	}
	appName := "iapp-{{.plat}}-{{.pool}}"
	flow.forward = func(c *check.C, env *Environment) {
		res := T("app-create", appName, "{{.plat}}", "-t", "{{.team}}", "-o", "{{.pool}}").Run(env)
		c.Assert(res, ResultOk)
		res = T("app-info", "-a", appName).Run(env)
		c.Assert(res, ResultOk)
		platRE := regexp.MustCompile(`(?s)Platform: (.*?)\n`)
		parts := platRE.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		lang := strings.Replace(parts[1], "iplat-", "", -1)
		res = T("app-deploy", "-a", appName, "{{.examplesdir}}/"+lang+"/").Run(env)
		c.Assert(res, ResultOk)
		res = T("app-info", "-a", appName).Run(env)
		c.Assert(res, ResultOk)
		addrRE := regexp.MustCompile(`(?s)Address: (.*?)\n`)
		parts = addrRE.FindStringSubmatch(res.Stdout.String())
		c.Assert(parts, check.HasLen, 2)
		cmd := NewCommand("curl", "-sSf", "http://"+parts[1])
		ok := retry(5*time.Minute, func() bool {
			res = cmd.Run(env)
			return res.ExitCode == 0
		})
		c.Assert(ok, check.Equals, true, check.Commentf("invalid result: %v", res))
	}
	flow.backward = func(c *check.C, env *Environment) {
		res := T("app-remove", "-y", "-a", appName).Run(env)
		c.Check(res, ResultOk)
	}
	return flow
}

func (s *S) TestBase(c *check.C) {
	env := NewEnvironment()
	if !env.Has("enabled") {
		return
	}
	var executedFlows []*ExecFlow
	defer func() {
		for i := len(executedFlows) - 1; i >= 0; i-- {
			executedFlows[i].Rollback(c, env)
		}
	}()
	for i := range flows {
		f := &flows[i]
		if len(f.provides) > 0 {
			providesAll := true
			for _, envVar := range f.provides {
				if env.Get(envVar) == "" {
					providesAll = false
					break
				}
			}
			if providesAll {
				continue
			}
		}
		executedFlows = append(executedFlows, f)
		f.Run(c, env)
	}
}
<|endoftext|>