<commit_before>package online

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"

	log "github.com/Sirupsen/logrus"
)

type Client struct {
	Token  string
	client *http.Client
}

func NewClient(token string) Client {
	c := Client{
		Token:  token,
		client: &http.Client{},
	}
	return c
}

func (c Client) getApiResource(resName string) ([]byte, error) {
	url := fmt.Sprintf("https://api.cloud.online.net/%s", resName)
	log.Debugf("url:%q", url)
	req, err := http.NewRequest("GET", url, nil)

	req.Header.Set("X-Auth-Token", c.Token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Status code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// Images return the list of available images
func (c Client) Images() ([]Image, error) {

	body, err := c.getApiResource("images")
	if err != nil {
		return nil, err
	}
	//log.Debug(string(body))

	var imgRes GetImagesResp
	err = json.Unmarshal(body, &imgRes)
	if err != nil {
		return nil, err
	}

	return imgRes.Images, nil
}

func (c Client) Servers() ([]Server, error) {

	body, err := c.getApiResource("servers")
	if err != nil {
		return nil, err
	}
	//log.Debug(string(body))

	var serRes GetServersResp
	err = json.Unmarshal(body, &serRes)
	if err != nil {
		return nil, err
	}

	return serRes.Servers, nil
}

func (c Client) Volumes() ([]Volume, error) {

	body, err := c.getApiResource("volumes")
	if err != nil {
		return nil, err
	}
	//log.Debug(string(body))

	var resp GetVolumesResp
	err = json.Unmarshal(body, &resp)
	if err != nil {
		return nil, err
	}

	return resp.Volumes, nil
}

func (c Client) Snapshots() ([]Snapshot, error) {

	body, err := c.getApiResource("snapshots")
	if err != nil {
		return nil, err
	}
	//log.Debug(string(body))

	var resp GetSnapshotsResp
	err = json.Unmarshal(body, &resp)
	if err != nil {
		return nil, err
	}

	return resp.Snapshots, nil
}
<commit_msg>Debug API respons when -v --verbose option used<commit_after>package online

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"

	log "github.com/Sirupsen/logrus"
)

type Client struct {
	Token  string
	client *http.Client
}

func NewClient(token string) Client {
	c := Client{
		Token:  token,
		client: &http.Client{},
	}
	return c
}

func (c Client) getApiResource(resName string) ([]byte, error) {
	url := fmt.Sprintf("https://api.cloud.online.net/%s", resName)
	log.Debugf("url:%q", url)
	req, err := http.NewRequest("GET", url, nil)

	req.Header.Set("X-Auth-Token", c.Token)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Status code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// Images return the list of available images
func (c Client) Images() ([]Image, error) {

	body, err := c.getApiResource("images")
	if err != nil {
		return nil, err
	}
	log.Debugf("API resp: %s", string(body))

	var imgRes GetImagesResp
	err = json.Unmarshal(body, &imgRes)
	if err != nil {
		return nil, err
	}

	return imgRes.Images, nil
}

func (c Client) Servers() ([]Server, error) {

	body, err := c.getApiResource("servers")
	if err != nil {
		return nil, err
	}
	log.Debugf("API resp: %s", string(body))

	var serRes GetServersResp
	err = json.Unmarshal(body, &serRes)
	if err != nil {
		return nil, err
	}

	return serRes.Servers, nil
}

func (c Client) Volumes() ([]Volume, error) {

	body, err := c.getApiResource("volumes")
	if err != nil {
		return nil, err
	}
	log.Debugf("API resp: %s", string(body))

	var resp GetVolumesResp
	err = json.Unmarshal(body, &resp)
	if err != nil {
		return nil, err
	}

	return resp.Volumes, nil
}

func (c Client) Snapshots() ([]Snapshot, error) {

	body, err := c.getApiResource("snapshots")
	if err != nil {
		return nil, err
	}
	log.Debugf("API resp: %s", string(body))

	var resp GetSnapshotsResp
	err = json.Unmarshal(body, &resp)
	if err != nil {
		return nil, err
	}

	return resp.Snapshots, nil
}
<|endoftext|>


<commit_before>package indexer

import (
	"encoding/json"
	"regexp"
	"strconv"
	"strings"
	// "fmt"
	"log"
	"path"
	// "strings"

	"gopkg.in/olivere/elastic.v3"
)

type ESIndexer struct {
	client *elastic.Client
}

type FileIndex struct {
	Blob     string     `json:"blob"`
	Metadata []Metadata `json:"metadata"`
	Content  string     `json:"content"`
}

type Metadata struct {
	Project string `json:"project"`
	Repo    string `json:"repo"`
	Refs    string `json:"refs"`
	Path    string `json:"path"`
	Ext     string `json:"ext"`
}

type Hit struct {
	Source Source `json:"_source"`
	// Highlight map[string][]string `json:"highlight"`
	Highlight []HighlightSource `json:"highlight"`
}
type Source struct {
	Blob     string     `json:"blob"`
	Metadata []Metadata `json:"metadata"`
}
type HighlightSource struct {
	Offset  int    `json:"offset"`
	Content string `json:"content"`
}

var LINE_TAG = regexp.MustCompile(`^\[([0-9]+)\]\s(.*)`)

func NewESIndexer() Indexer {
	client, err := elastic.NewClient(elastic.SetURL())
	if err != nil {
		panic(err)
	}
	i := &ESIndexer{client: client}
	i.Init()
	return i
}

func (esi *ESIndexer) Init() {

	// esi.client.DeleteIndex("gosource").Do()
	_, err := esi.client.CreateIndex("gosource").BodyString(`
{
	settings: {
		analysis: {
			filter: {
				pos_filter: {
					type: "kuromoji_part_of_speech",
					stoptags: ["助詞-格助詞-一般", "助詞-終助詞"]
				},
				greek_lowercase_filter: {
					type: "lowercase",
					language: "greek"
				}
			},
			char_filter: {
				remove_tags: {
					type: "pattern_replace",
					pattern: "^\\[[0-9]+\\]\\u0020",
					flags: "MULTILINE",
					replacement: ""
				}
			},
			analyzer: {
				path_analyzer: {
					type: "custom",
					tokenizer: "path_tokenizer"
				},
				kuromoji_analyzer: {
					type: "custom",
					tokenizer: "kuromoji_tokenizer",
					char_filter: ["remove_tags"],
					filter: ["kuromoji_baseform", "pos_filter", "greek_lowercase_filter", "cjk_width"]
				}
			},
			tokenizer: {
				path_tokenizer: {
					type: "path_hierarchy",
					reverse: true
				}
			}
		}
	},
	mappings: {
		file: {
			properties: {
				blob: {
					type: "string",
					index: "not_analyzed"
				},
				metadata: {
					type: "nested",
					properties: {
						project: {
							type: "multi_field",
							fields: {
								project: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						repository: {
							type: "multi_field",
							fields: {
								repository: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						refs: {
							type: "multi_field",
							fields: {
								refs: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						path: {
							type: "string",
							analyzer: "path_analyzer"
						},
						ext: {
							type: "string",
							index: "not_analyzed"
						}
					}
				},
				contents: {
					type: "string",
					index_options: "offsets",
					analyzer: "kuromoji_analyzer"
				}
			}
		}
	}
}
		`).Do()

	if err != nil {
		log.Println(err)
	}
}

func (esi *ESIndexer) CreateFileIndex(project string, repo string, branch string, filePath string, blob string, content string) error {

	ext := path.Ext(filePath)

	fileIndex := FileIndex{Blob: blob, Metadata: []Metadata{Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext}}, Content: content}

	_, err := esi.client.Index().
		Index("gosource").
		Type("file").
		Id(blob).
		BodyJson(fileIndex).
		Refresh(true).
		Do()

	if err != nil {
		return err
	}
	return nil
}

func (esi *ESIndexer) UpsertFileIndex(project string, repo string, branch string, filePath string, blob string, content string) error {

	ext := path.Ext(filePath)

	get, err := esi.client.Get().
		Index("gosource").
		Type("file").
		Id(blob).
		Do()

	if err == nil && get.Found {
		var fileIndex FileIndex
		if err := json.Unmarshal(*get.Source, &fileIndex); err != nil {
			return err
		}
		f := func(x Metadata, i int) bool {
			return x.Project == project &&
				x.Repo == repo &&
				x.Refs == branch &&
				x.Path == filePath
		}
		found := find(f, fileIndex.Metadata)
		if found != nil {
			fileIndex.Metadata = append(fileIndex.Metadata, Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext})
		}

		_, err := esi.client.Update().
			Index("gosource").
			Type("file").
			Id(blob).
			Doc(fileIndex).
			Do()

		if err != nil {
			log.Println("Upsert Doc error", err)
			return err
		}

	} else {
		lines := strings.Split(content, "\n")
		newLines := []string{}
		for i, l := range lines {
			newLines = append(newLines, "["+strconv.Itoa(i+1)+"] "+l)
		}

		fileIndex := FileIndex{Blob: blob, Metadata: []Metadata{Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext}}, Content: strings.Join(newLines, "\n")}

		_, err := esi.client.Index().
			Index("gosource").
			Type("file").
			Id(blob).
			BodyJson(fileIndex).
			Refresh(true).
			Do()

		if err != nil {
			log.Println("Add Doc error", err)
			return err
		}
	}

	log.Println("Indexed!")
	return nil
}

func (esi *ESIndexer) SearchQuery(query string) []Hit {
	// termQuery := elastic.NewTermsQuery("content", strings.Split(query, " "))
	q := elastic.NewQueryStringQuery(query).DefaultField("content").DefaultOperator("AND")
	searchResult, err := esi.client.Search().
		Index("gosource"). // search in index "twitter"
		FetchSourceContext(elastic.NewFetchSourceContext(true).Include("blob", "metadata")).
		Query(q). // specify the query
		Highlight(elastic.NewHighlight().Field("content").PreTags("@GITK_MARK_PRE@").PostTags("@GITK_MARK_POST@")).
		Sort("metadata.path", true). // sort by "user" field, ascending
		From(0).Size(10).            // take documents 0-9
		Pretty(true).                // pretty print request and response JSON
		Do()                         // execute

	if err != nil {
		log.Println("error", err)
		return []Hit{}
	}

	list := []Hit{}
	if searchResult.Hits.TotalHits > 0 {
		for _, hit := range searchResult.Hits.Hits {
			// hit.Index contains the name of the index

			// Deserialize hit.Source into a Tweet (could also be just a map[string]interface{}).
			var s Source
			json.Unmarshal(*hit.Source, &s)

			hsList := []HighlightSource{}
			for _, hc := range hit.Highlight["content"] {
				list := []string{}
				first := 0
				for _, l := range strings.Split(hc, "\n") {
					groups := LINE_TAG.FindAllStringSubmatch(l, 1)
					if len(groups) == 1 {
						if first == 0 {
							if strings.TrimSpace(groups[0][2]) != "" {
								first, _ = strconv.Atoi(groups[0][1])
								list = append(list, groups[0][2])
							}
						} else {
							list = append(list, groups[0][2])
						}
					}
				}
				hs := HighlightSource{Offset: first, Content: strings.Join(list, "\n")}
				hsList = append(hsList, hs)
			}

			h := Hit{Source: s, Highlight: hsList}
			list = append(list, h)
		}
	}
	return list
}

func find(f func(s Metadata, i int) bool, s []Metadata) *Metadata {
	for index, x := range s {
		if f(x, index) == true {
			return &x
		}
	}
	return nil
}

func filter(f func(s Metadata, i int) bool, s []Metadata) []Metadata {
	ans := make([]Metadata, 0)
	for index, x := range s {
		if f(x, index) == true {
			ans = append(ans, x)
		}
	}
	return ans
}
<commit_msg>Fix metadata bug<commit_after>package indexer

import (
	"encoding/json"
	"regexp"
	"strconv"
	"strings"
	// "fmt"
	"log"
	"path"
	// "strings"

	"gopkg.in/olivere/elastic.v3"
)

type ESIndexer struct {
	client *elastic.Client
}

type FileIndex struct {
	Blob     string     `json:"blob"`
	Metadata []Metadata `json:"metadata"`
	Content  string     `json:"content"`
}

type Metadata struct {
	Project string `json:"project"`
	Repo    string `json:"repo"`
	Refs    string `json:"refs"`
	Path    string `json:"path"`
	Ext     string `json:"ext"`
}

type Hit struct {
	Source Source `json:"_source"`
	// Highlight map[string][]string `json:"highlight"`
	Highlight []HighlightSource `json:"highlight"`
}
type Source struct {
	Blob     string     `json:"blob"`
	Metadata []Metadata `json:"metadata"`
}
type HighlightSource struct {
	Offset  int    `json:"offset"`
	Content string `json:"content"`
}

var LINE_TAG = regexp.MustCompile(`^\[([0-9]+)\]\s(.*)`)

func NewESIndexer() Indexer {
	client, err := elastic.NewClient(elastic.SetURL())
	if err != nil {
		panic(err)
	}
	i := &ESIndexer{client: client}
	i.Init()
	return i
}

func (esi *ESIndexer) Init() {

	// esi.client.DeleteIndex("gosource").Do()
	_, err := esi.client.CreateIndex("gosource").BodyString(`
{
	settings: {
		analysis: {
			filter: {
				pos_filter: {
					type: "kuromoji_part_of_speech",
					stoptags: ["助詞-格助詞-一般", "助詞-終助詞"]
				},
				greek_lowercase_filter: {
					type: "lowercase",
					language: "greek"
				}
			},
			char_filter: {
				remove_tags: {
					type: "pattern_replace",
					pattern: "^\\[[0-9]+\\]\\u0020",
					flags: "MULTILINE",
					replacement: ""
				}
			},
			analyzer: {
				path_analyzer: {
					type: "custom",
					tokenizer: "path_tokenizer"
				},
				kuromoji_analyzer: {
					type: "custom",
					tokenizer: "kuromoji_tokenizer",
					char_filter: ["remove_tags"],
					filter: ["kuromoji_baseform", "pos_filter", "greek_lowercase_filter", "cjk_width"]
				}
			},
			tokenizer: {
				path_tokenizer: {
					type: "path_hierarchy",
					reverse: true
				}
			}
		}
	},
	mappings: {
		file: {
			properties: {
				blob: {
					type: "string",
					index: "not_analyzed"
				},
				metadata: {
					type: "nested",
					properties: {
						project: {
							type: "multi_field",
							fields: {
								project: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						repository: {
							type: "multi_field",
							fields: {
								repository: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						refs: {
							type: "multi_field",
							fields: {
								refs: {
									type: "string",
									index: "analyzed"
								},
								full: {
									type: "string",
									index: "not_analyzed"
								}
							}
						},
						path: {
							type: "string",
							analyzer: "path_analyzer"
						},
						ext: {
							type: "string",
							index: "not_analyzed"
						}
					}
				},
				contents: {
					type: "string",
					index_options: "offsets",
					analyzer: "kuromoji_analyzer"
				}
			}
		}
	}
}
		`).Do()

	if err != nil {
		log.Println(err)
	}
}

func (esi *ESIndexer) CreateFileIndex(project string, repo string, branch string, filePath string, blob string, content string) error {

	ext := path.Ext(filePath)

	fileIndex := FileIndex{Blob: blob, Metadata: []Metadata{Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext}}, Content: content}

	_, err := esi.client.Index().
		Index("gosource").
		Type("file").
		Id(blob).
		BodyJson(fileIndex).
		Refresh(true).
		Do()

	if err != nil {
		return err
	}
	return nil
}

func (esi *ESIndexer) UpsertFileIndex(project string, repo string, branch string, filePath string, blob string, content string) error {

	ext := path.Ext(filePath)

	get, err := esi.client.Get().
		Index("gosource").
		Type("file").
		Id(blob).
		Do()

	if err == nil && get.Found {
		var fileIndex FileIndex
		if err := json.Unmarshal(*get.Source, &fileIndex); err != nil {
			return err
		}
		f := func(x Metadata, i int) bool {
			return x.Project == project &&
				x.Repo == repo &&
				x.Refs == branch &&
				x.Path == filePath
		}
		found := find(f, fileIndex.Metadata)
		if found == nil {
			fileIndex.Metadata = append(fileIndex.Metadata, Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext})
		}

		_, err := esi.client.Update().
			Index("gosource").
			Type("file").
			Id(blob).
			Doc(fileIndex).
			Do()

		if err != nil {
			log.Println("Upsert Doc error", err)
			return err
		}

	} else {
		lines := strings.Split(content, "\n")
		newLines := []string{}
		for i, l := range lines {
			newLines = append(newLines, "["+strconv.Itoa(i+1)+"] "+l)
		}

		fileIndex := FileIndex{Blob: blob, Metadata: []Metadata{Metadata{Project: project, Repo: repo, Refs: branch, Path: filePath, Ext: ext}}, Content: strings.Join(newLines, "\n")}

		_, err := esi.client.Index().
			Index("gosource").
			Type("file").
			Id(blob).
			BodyJson(fileIndex).
			Refresh(true).
			Do()

		if err != nil {
			log.Println("Add Doc error", err)
			return err
		}
	}

	log.Println("Indexed!")
	return nil
}

func (esi *ESIndexer) SearchQuery(query string) []Hit {
	// termQuery := elastic.NewTermsQuery("content", strings.Split(query, " "))
	q := elastic.NewQueryStringQuery(query).DefaultField("content").DefaultOperator("AND")
	searchResult, err := esi.client.Search().
		Index("gosource"). // search in index "twitter"
		FetchSourceContext(elastic.NewFetchSourceContext(true).Include("blob", "metadata")).
		Query(q). // specify the query
		Highlight(elastic.NewHighlight().Field("content").PreTags("@GITK_MARK_PRE@").PostTags("@GITK_MARK_POST@")).
		Sort("metadata.path", true). // sort by "user" field, ascending
		From(0).Size(10).            // take documents 0-9
		Pretty(true).                // pretty print request and response JSON
		Do()                         // execute

	if err != nil {
		log.Println("error", err)
		return []Hit{}
	}

	list := []Hit{}
	if searchResult.Hits.TotalHits > 0 {
		for _, hit := range searchResult.Hits.Hits {
			// hit.Index contains the name of the index

			// Deserialize hit.Source into a Tweet (could also be just a map[string]interface{}).
			var s Source
			json.Unmarshal(*hit.Source, &s)

			hsList := []HighlightSource{}
			for _, hc := range hit.Highlight["content"] {
				list := []string{}
				first := 0
				for _, l := range strings.Split(hc, "\n") {
					groups := LINE_TAG.FindAllStringSubmatch(l, 1)
					if len(groups) == 1 {
						if first == 0 {
							if strings.TrimSpace(groups[0][2]) != "" {
								first, _ = strconv.Atoi(groups[0][1])
								list = append(list, groups[0][2])
							}
						} else {
							list = append(list, groups[0][2])
						}
					}
				}
				hs := HighlightSource{Offset: first, Content: strings.Join(list, "\n")}
				hsList = append(hsList, hs)
			}

			h := Hit{Source: s, Highlight: hsList}
			list = append(list, h)
		}
	}
	return list
}

func find(f func(s Metadata, i int) bool, s []Metadata) *Metadata {
	for index, x := range s {
		if f(x, index) == true {
			return &x
		}
	}
	return nil
}

func filter(f func(s Metadata, i int) bool, s []Metadata) []Metadata {
	ans := make([]Metadata, 0)
	for index, x := range s {
		if f(x, index) == true {
			ans = append(ans, x)
		}
	}
	return ans
}
<|endoftext|>


<commit_before>package lxd

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

// Event handling functions

// getEvents connects to the LXD monitoring interface
func (r *ProtocolLXD) getEvents(allProjects bool) (*EventListener, error) {
	// Prevent anything else from interacting with the listeners
	r.eventListenersLock.Lock()
	defer r.eventListenersLock.Unlock()

	ctx, cancel := context.WithCancel(context.Background())

	// Setup a new listener
	listener := EventListener{
		r:         r,
		ctx:       ctx,
		ctxCancel: cancel,
	}

	connInfo, _ := r.GetConnectionInfo()
	if connInfo.Project == "" {
		return nil, fmt.Errorf("Unexpected empty project in connection info")
	}

	if !allProjects {
		listener.projectName = connInfo.Project
	}

	// There is an existing Go routine for the required project filter, so just add another target.
	if r.eventListeners[listener.projectName] != nil {
		r.eventListeners[listener.projectName] = append(r.eventListeners[listener.projectName], &listener)
		return &listener, nil
	}

	// Setup a new connection with LXD
	var url string
	var err error
	if allProjects {
		url, err = r.setQueryAttributes("/events?all-projects=true")
	} else {
		url, err = r.setQueryAttributes("/events")
	}
	if err != nil {
		return nil, err
	}

	// Connect websocket and save.
	wsConn, err := r.websocket(url)
	if err != nil {
		return nil, err
	}

	r.eventConns[listener.projectName] = wsConn // Save for others to use.

	// Initialize the event listener list if we were able to connect to the events websocket.
	r.eventListeners[listener.projectName] = []*EventListener{&listener}

	// Spawn a watcher that will close the websocket connection after all
	// listeners are gone.
	stopCh := make(chan struct{})
	go func() {
		for {
			select {
			case <-time.After(time.Minute):
			case <-r.chConnected:
			case <-stopCh:
			}

			r.eventListenersLock.Lock()
			if len(r.eventListeners[listener.projectName]) == 0 {
				// We don't need the connection anymore, disconnect and clear.
				if r.eventListeners[listener.projectName] != nil {
					r.eventConns[listener.projectName].Close()
					r.eventConns[listener.projectName] = nil
				}

				r.eventListeners[listener.projectName] = nil
				r.eventListenersLock.Unlock()

				return
			}
			r.eventListenersLock.Unlock()
		}
	}()

	// Spawn the listener
	go func() {
		for {
			_, data, err := wsConn.ReadMessage()
			if err != nil {
				// Prevent anything else from interacting with the listeners
				r.eventListenersLock.Lock()
				defer r.eventListenersLock.Unlock()

				// Tell all the current listeners about the failure
				for _, listener := range r.eventListeners[listener.projectName] {
					listener.err = err
					listener.ctxCancel()
				}

				// And remove them all from the list so that when watcher routine runs it will
				// close the websocket connection.
				r.eventListeners[listener.projectName] = nil

				close(stopCh) // Instruct watcher go routine to cleanup.

				return
			}

			// Attempt to unpack the message
			event := api.Event{}
			err = json.Unmarshal(data, &event)
			if err != nil {
				continue
			}

			// Extract the message type
			if event.Type == "" {
				continue
			}

			// Send the message to all handlers
			r.eventListenersLock.Lock()
			for _, listener := range r.eventListeners[listener.projectName] {
				listener.targetsLock.Lock()
				for _, target := range listener.targets {
					if target.types != nil && !shared.StringInSlice(event.Type, target.types) {
						continue
					}

					go target.function(event)
				}
				listener.targetsLock.Unlock()
			}
			r.eventListenersLock.Unlock()
		}
	}()

	return &listener, nil
}

// GetEvents gets the events for the project defined on the client.
func (r *ProtocolLXD) GetEvents() (*EventListener, error) {
	return r.getEvents(false)
}

// GetEventsAllProjects gets events for all projects.
func (r *ProtocolLXD) GetEventsAllProjects() (*EventListener, error) {
	return r.getEvents(true)
}
<commit_msg>client/lxd/events: Use r.eventConnsLock to control access to r.eventConns<commit_after>package lxd

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/gorilla/websocket"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

// Event handling functions

// getEvents connects to the LXD monitoring interface
func (r *ProtocolLXD) getEvents(allProjects bool) (*EventListener, error) {
	// Prevent anything else from interacting with the listeners
	r.eventListenersLock.Lock()
	defer r.eventListenersLock.Unlock()

	ctx, cancel := context.WithCancel(context.Background())

	// Setup a new listener
	listener := EventListener{
		r:         r,
		ctx:       ctx,
		ctxCancel: cancel,
	}

	connInfo, _ := r.GetConnectionInfo()
	if connInfo.Project == "" {
		return nil, fmt.Errorf("Unexpected empty project in connection info")
	}

	if !allProjects {
		listener.projectName = connInfo.Project
	}

	// There is an existing Go routine for the required project filter, so just add another target.
	if r.eventListeners[listener.projectName] != nil {
		r.eventListeners[listener.projectName] = append(r.eventListeners[listener.projectName], &listener)
		return &listener, nil
	}

	// Setup a new connection with LXD
	var url string
	var err error
	if allProjects {
		url, err = r.setQueryAttributes("/events?all-projects=true")
	} else {
		url, err = r.setQueryAttributes("/events")
	}
	if err != nil {
		return nil, err
	}

	// Connect websocket and save.
	wsConn, err := r.websocket(url)
	if err != nil {
		return nil, err
	}

	r.eventConnsLock.Lock()
	r.eventConns[listener.projectName] = wsConn // Save for others to use.
	r.eventConnsLock.Unlock()

	// Initialize the event listener list if we were able to connect to the events websocket.
	r.eventListeners[listener.projectName] = []*EventListener{&listener}

	// Spawn a watcher that will close the websocket connection after all
	// listeners are gone.
	stopCh := make(chan struct{})
	go func() {
		for {
			select {
			case <-time.After(time.Minute):
			case <-r.chConnected:
			case <-stopCh:
			}

			r.eventListenersLock.Lock()
			r.eventConnsLock.Lock()
			if len(r.eventListeners[listener.projectName]) == 0 {
				// We don't need the connection anymore, disconnect and clear.
				if r.eventListeners[listener.projectName] != nil {
					r.eventConns[listener.projectName].Close()
					delete(r.eventConns, listener.projectName)
				}

				r.eventListeners[listener.projectName] = nil
				r.eventListenersLock.Unlock()
				r.eventConnsLock.Unlock()

				return
			}
			r.eventListenersLock.Unlock()
			r.eventConnsLock.Unlock()
		}
	}()

	// Spawn the listener
	go func() {
		for {
			_, data, err := wsConn.ReadMessage()
			if err != nil {
				// Prevent anything else from interacting with the listeners
				r.eventListenersLock.Lock()
				defer r.eventListenersLock.Unlock()

				// Tell all the current listeners about the failure
				for _, listener := range r.eventListeners[listener.projectName] {
					listener.err = err
					listener.ctxCancel()
				}

				// And remove them all from the list so that when watcher routine runs it will
				// close the websocket connection.
				r.eventListeners[listener.projectName] = nil

				close(stopCh) // Instruct watcher go routine to cleanup.

				return
			}

			// Attempt to unpack the message
			event := api.Event{}
			err = json.Unmarshal(data, &event)
			if err != nil {
				continue
			}

			// Extract the message type
			if event.Type == "" {
				continue
			}

			// Send the message to all handlers
			r.eventListenersLock.Lock()
			for _, listener := range r.eventListeners[listener.projectName] {
				listener.targetsLock.Lock()
				for _, target := range listener.targets {
					if target.types != nil && !shared.StringInSlice(event.Type, target.types) {
						continue
					}

					go target.function(event)
				}
				listener.targetsLock.Unlock()
			}
			r.eventListenersLock.Unlock()
		}
	}()

	return &listener, nil
}

// GetEvents gets the events for the project defined on the client.
func (r *ProtocolLXD) GetEvents() (*EventListener, error) {
	return r.getEvents(false)
}

// GetEventsAllProjects gets events for all projects.
func (r *ProtocolLXD) GetEventsAllProjects() (*EventListener, error) {
	return r.getEvents(true)
}
<|endoftext|>


<commit_before>package lxd

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/api"
)

// The Operation type represents an ongoing LXD operation (asynchronous processing)
type operation struct {
	api.Operation

	r            *ProtocolLXD
	listener     *EventListener
	handlerReady bool
	handlerLock  sync.Mutex

	chActive chan bool
}

// AddHandler adds a function to be called whenever an event is received
func (op *operation) AddHandler(function func(api.Operation)) (*EventTarget, error) {
	// Make sure we have a listener setup
	err := op.setupListener()
	if err != nil {
		return nil, err
	}

	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// If we're done already, just return
	if op.StatusCode.IsFinal() {
		return nil, nil
	}

	// Wrap the function to filter unwanted messages
	wrapped := func(event api.Event) {
		newOp := api.Operation{}
		err := json.Unmarshal(event.Metadata, &newOp)
		if err != nil || newOp.ID != op.ID {
			return
		}

		function(newOp)
	}

	return op.listener.AddHandler([]string{"operation"}, wrapped)
}

// Cancel will request that LXD cancels the operation (if supported)
func (op *operation) Cancel() error {
	return op.r.DeleteOperation(op.ID)
}

// Get returns the API operation struct
func (op *operation) Get() api.Operation {
	return op.Operation
}

// GetWebsocket returns a raw websocket connection from the operation
func (op *operation) GetWebsocket(secret string) (*websocket.Conn, error) {
	return op.r.GetOperationWebsocket(op.ID, secret)
}

// RemoveHandler removes a function to be called whenever an event is received
func (op *operation) RemoveHandler(target *EventTarget) error {
	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// If the listener is gone, just return
	if op.listener == nil {
		return nil
	}

	return op.listener.RemoveHandler(target)
}

// Refresh pulls the current version of the operation and updates the struct
func (op *operation) Refresh() error {
	// Get the current version of the operation
	newOp, _, err := op.r.GetOperation(op.ID)
	if err != nil {
		return err
	}

	// Update the operation struct
	op.Operation = *newOp

	return nil
}

// Wait lets you wait until the operation reaches a final state
func (op *operation) Wait() error {
	// Check if not done already
	if op.StatusCode.IsFinal() {
		if op.Err != "" {
			return fmt.Errorf(op.Err)
		}

		return nil
	}

	// Make sure we have a listener setup
	err := op.setupListener()
	if err != nil {
		return err
	}

	<-op.chActive

	// We're done, parse the result
	if op.Err != "" {
		return fmt.Errorf(op.Err)
	}

	return nil
}

func (op *operation) setupListener() error {
	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// We already have a listener setup
	if op.handlerReady {
		return nil
	}
	op.handlerReady = true

	// Get a new listener
	if op.listener == nil {
		listener, err := op.r.GetEvents()
		if err != nil {
			return err
		}

		op.listener = listener
	}

	// Setup the handler
	chReady := make(chan bool)
	_, err := op.listener.AddHandler([]string{"operation"}, func(event api.Event) {
		<-chReady

		// We don't want concurrency while processing events
		op.handlerLock.Lock()
		defer op.handlerLock.Unlock()

		// Check if we're done already (because of another event)
		if op.listener == nil {
			return
		}

		// Get an operation struct out of this data
		newOp := api.Operation{}
		err := json.Unmarshal(event.Metadata, &newOp)
		if err != nil || newOp.ID != op.ID {
			return
		}

		// Update the struct
		op.Operation = newOp

		// And check if we're done
		if op.StatusCode.IsFinal() {
			op.listener.Disconnect()
			op.listener = nil
			close(op.chActive)
			return
		}
	})
	if err != nil {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		return err
	}

	// Monitor event listener
	go func() {
		<-chReady

		// We don't want concurrency while accessing the listener
		op.handlerLock.Lock()

		// Check if we're done already (because of another event)
		listener := op.listener
		if listener == nil {
			op.handlerLock.Unlock()
			return
		}
		op.handlerLock.Unlock()

		// Wait for the listener or operation to be done
		select {
		case <-listener.chActive:
			op.handlerLock.Lock()
			if op.listener != nil {
				op.Err = fmt.Sprintf("%v", listener.err)
				close(op.chActive)
			}
			op.handlerLock.Unlock()
		case <-op.chActive:
			return
		}
	}()

	// And do a manual refresh to avoid races
	err = op.Refresh()
	if err != nil {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		return err
	}

	// Check if not done already
	if op.StatusCode.IsFinal() {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		if op.Err != "" {
			return fmt.Errorf(op.Err)
		}

		return nil
	}

	// Start processing background updates
	close(chReady)

	return nil
}

// The remoteOperation type represents an ongoing LXD operation between two servers
type remoteOperation struct {
	targetOp Operation

	handlers []func(api.Operation)

	chDone chan bool
	chPost chan bool
	err    error
}

// AddHandler adds a function to be called whenever an event is received
func (op *remoteOperation) AddHandler(function func(api.Operation)) (*EventTarget, error) {
	var err error
	var target *EventTarget

	// Attach to the existing target operation
	if op.targetOp != nil {
		target, err = op.targetOp.AddHandler(function)
		if err != nil {
			return nil, err
		}
	} else {
		// Generate a mock EventTarget
		target = &EventTarget{
			function: func(api.Event) { function(api.Operation{}) },
			types:    []string{"operation"},
		}
	}

	// Add the handler to our list
	op.handlers = append(op.handlers, function)

	return target, nil
}

// CancelTarget attempts to cancel the target operation
func (op *remoteOperation) CancelTarget() error {
	if op.targetOp == nil {
		return fmt.Errorf("No associated target operation")
	}

	return op.targetOp.Cancel()
}

// GetTarget returns the target operation
func (op *remoteOperation) GetTarget() (*api.Operation, error) {
	if op.targetOp == nil {
		return nil, fmt.Errorf("No associated target operation")
	}

	opAPI := op.targetOp.Get()
	return &opAPI, nil
}

// Wait lets you wait until the operation reaches a final state
func (op *remoteOperation) Wait() error {
	<-op.chDone

	if op.chPost != nil {
		<-op.chPost
	}

	return op.err
}
<commit_msg>client/operations: Fixes race conditions<commit_after>package lxd

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/api"
)

// The Operation type represents an ongoing LXD operation (asynchronous processing)
type operation struct {
	api.Operation

	r            *ProtocolLXD
	listener     *EventListener
	handlerReady bool
	handlerLock  sync.Mutex

	chActive chan bool
}

// AddHandler adds a function to be called whenever an event is received
func (op *operation) AddHandler(function func(api.Operation)) (*EventTarget, error) {
	// Make sure we have a listener setup
	err := op.setupListener()
	if err != nil {
		return nil, err
	}

	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// If we're done already, just return
	if op.StatusCode.IsFinal() {
		return nil, nil
	}

	// Wrap the function to filter unwanted messages
	wrapped := func(event api.Event) {
		op.handlerLock.Lock()

		newOp := api.Operation{}
		err := json.Unmarshal(event.Metadata, &newOp)
		if err != nil || newOp.ID != op.ID {
			op.handlerLock.Unlock()
			return
		}
		op.handlerLock.Unlock()

		function(newOp)
	}

	return op.listener.AddHandler([]string{"operation"}, wrapped)
}

// Cancel will request that LXD cancels the operation (if supported)
func (op *operation) Cancel() error {
	return op.r.DeleteOperation(op.ID)
}

// Get returns the API operation struct
func (op *operation) Get() api.Operation {
	return op.Operation
}

// GetWebsocket returns a raw websocket connection from the operation
func (op *operation) GetWebsocket(secret string) (*websocket.Conn, error) {
	return op.r.GetOperationWebsocket(op.ID, secret)
}

// RemoveHandler removes a function to be called whenever an event is received
func (op *operation) RemoveHandler(target *EventTarget) error {
	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// If the listener is gone, just return
	if op.listener == nil {
		return nil
	}

	return op.listener.RemoveHandler(target)
}

// Refresh pulls the current version of the operation and updates the struct
func (op *operation) Refresh() error {
	// Get the current version of the operation
	newOp, _, err := op.r.GetOperation(op.ID)
	if err != nil {
		return err
	}

	// Update the operation struct
	op.Operation = *newOp

	return nil
}

// Wait lets you wait until the operation reaches a final state
func (op *operation) Wait() error {
	op.handlerLock.Lock()
	// Check if not done already
	if op.StatusCode.IsFinal() {
		if op.Err != "" {
			op.handlerLock.Unlock()
			return fmt.Errorf(op.Err)
		}

		op.handlerLock.Unlock()
		return nil
	}
	op.handlerLock.Unlock()

	// Make sure we have a listener setup
	err := op.setupListener()
	if err != nil {
		return err
	}

	<-op.chActive

	// We're done, parse the result
	if op.Err != "" {
		return fmt.Errorf(op.Err)
	}

	return nil
}

func (op *operation) setupListener() error {
	// Make sure we're not racing with ourselves
	op.handlerLock.Lock()
	defer op.handlerLock.Unlock()

	// We already have a listener setup
	if op.handlerReady {
		return nil
	}
	op.handlerReady = true

	// Get a new listener
	if op.listener == nil {
		listener, err := op.r.GetEvents()
		if err != nil {
			return err
		}

		op.listener = listener
	}

	// Setup the handler
	chReady := make(chan bool)
	_, err := op.listener.AddHandler([]string{"operation"}, func(event api.Event) {
		<-chReady

		// We don't want concurrency while processing events
		op.handlerLock.Lock()
		defer op.handlerLock.Unlock()

		// Check if we're done already (because of another event)
		if op.listener == nil {
			return
		}

		// Get an operation struct out of this data
		newOp := api.Operation{}
		err := json.Unmarshal(event.Metadata, &newOp)
		if err != nil || newOp.ID != op.ID {
			return
		}

		// Update the struct
		op.Operation = newOp

		// And check if we're done
		if op.StatusCode.IsFinal() {
			op.listener.Disconnect()
			op.listener = nil
			close(op.chActive)
			return
		}
	})
	if err != nil {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		return err
	}

	// Monitor event listener
	go func() {
		<-chReady

		// We don't want concurrency while accessing the listener
		op.handlerLock.Lock()

		// Check if we're done already (because of another event)
		listener := op.listener
		if listener == nil {
			op.handlerLock.Unlock()
			return
		}
		op.handlerLock.Unlock()

		// Wait for the listener or operation to be done
		select {
		case <-listener.chActive:
			op.handlerLock.Lock()
			if op.listener != nil {
				op.Err = fmt.Sprintf("%v", listener.err)
				close(op.chActive)
			}
			op.handlerLock.Unlock()
		case <-op.chActive:
			return
		}
	}()

	// And do a manual refresh to avoid races
	err = op.Refresh()
	if err != nil {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		return err
	}

	// Check if not done already
	if op.StatusCode.IsFinal() {
		op.listener.Disconnect()
		op.listener = nil
		close(op.chActive)
		close(chReady)

		if op.Err != "" {
			return fmt.Errorf(op.Err)
		}

		return nil
	}

	// Start processing background updates
	close(chReady)

	return nil
}

// The remoteOperation type represents an ongoing LXD operation between two servers
type remoteOperation struct {
	targetOp Operation

	handlers []func(api.Operation)

	chDone chan bool
	chPost chan bool
	err    error
}

// AddHandler adds a function to be called whenever an event is received
func (op *remoteOperation) AddHandler(function func(api.Operation)) (*EventTarget, error) {
	var err error
	var target *EventTarget

	// Attach to the existing target operation
	if op.targetOp != nil {
		target, err = op.targetOp.AddHandler(function)
		if err != nil {
			return nil, err
		}
	} else {
		// Generate a mock EventTarget
		target = &EventTarget{
			function: func(api.Event) { function(api.Operation{}) },
			types:    []string{"operation"},
		}
	}

	// Add the handler to our list
	op.handlers = append(op.handlers, function)

	return target, nil
}

// CancelTarget attempts to cancel the target operation
func (op *remoteOperation) CancelTarget() error {
	if op.targetOp == nil {
		return fmt.Errorf("No associated target operation")
	}

	return op.targetOp.Cancel()
}

// GetTarget returns the target operation
func (op *remoteOperation) GetTarget() (*api.Operation, error) {
	if op.targetOp == nil {
		return nil, fmt.Errorf("No associated target operation")
	}

	opAPI := op.targetOp.Get()
	return &opAPI, nil
}

// Wait lets you wait until the operation reaches a final state
func (op *remoteOperation) Wait() error {
	<-op.chDone

	if op.chPost != nil {
		<-op.chPost
	}

	return op.err
}
<|endoftext|>


<commit_before>package stage

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/flant/dapp/pkg/config"
	"github.com/flant/dapp/pkg/image"
	"github.com/flant/dapp/pkg/slug"
	"github.com/flant/dapp/pkg/util"
)

type StageName string

const (
	From                        StageName = "from"
	BeforeInstall               StageName = "before_install"
	ArtifactImportBeforeInstall StageName = "before_install_artifact"
	GAArchive                   StageName = "g_a_archive"
	Install                     StageName = "install"
	ArtifactImportAfterInstall  StageName = "after_install_artifact"
	BeforeSetup                 StageName = "before_setup"
	ArtifactImportBeforeSetup   StageName = "before_setup_artifact"
	Setup                       StageName = "setup"
	ArtifactImportAfterSetup    StageName = "after_setup_artifact"
	GAPostSetupPatch            StageName = "g_a_post_setup_patch"
	GALatestPatch               StageName = "g_a_latest_patch"
	DockerInstructions          StageName = "docker_instructions"
)

const (
	mountTmpDirLabel          = "dapp-mount-type-tmp-dir"
	mountBuildDIrLabel        = "dapp-mount-type-build-dir"
	mountCustomDirLabelPrefix = "dapp-mount-type-custom-dir-"
)

type NewBaseStageOptions struct {
	DimgName         string
	ConfigMounts     []*config.Mount
	DimgTmpDir       string
	ContainerDappDir string
	ProjectBuildDir  string
}

func newBaseStage(name StageName, options *NewBaseStageOptions) *BaseStage {
	s := &BaseStage{}
	s.name = name
	s.dimgName = options.DimgName
	s.configMounts = options.ConfigMounts
	s.projectBuildDir = options.ProjectBuildDir
	s.dimgTmpDir = options.DimgTmpDir
	s.containerDappDir = options.ContainerDappDir
	return s
}

type BaseStage struct {
	name             StageName
	dimgName         string
	signature        string
	image            image.Image
	gitArtifacts     []*GitArtifact
	dimgTmpDir       string
	containerDappDir string
	projectBuildDir  string
	configMounts     []*config.Mount
}

func (s *BaseStage) Name() StageName {
	if s.name != "" {
		return s.name
	}

	panic("name must be defined!")
}

func (s *BaseStage) GetDependencies(_ Conveyor, _ image.Image) (string, error) {
	panic("method must be implemented!")
}

func (s *BaseStage) IsEmpty(_ Conveyor, _ image.Image) (bool, error) {
	return false, nil
}

func (s *BaseStage) ShouldBeReset(_ image.Image) (bool, error) {
	return false, nil
}

func (s *BaseStage) PrepareImage(_ Conveyor, prevBuiltImage, image image.Image) error {
	/*
	 * NOTE: BaseStage.PrepareImage does not called in From.PrepareImage.
	 * NOTE: Take into account when adding new base PrepareImage steps.
	 */

	serviceMounts := s.getServiceMounts(prevBuiltImage)
	s.addServiceMountsLabels(serviceMounts, image)
	if err := s.addServiceMountsVolumes(serviceMounts, image); err != nil {
		return fmt.Errorf("error adding mounts volumes: %s", err)
	}

	customMounts := s.getCustomMounts(prevBuiltImage)
	s.addCustomMountLabels(customMounts, image)
	if err := s.addCustomMountVolumes(customMounts, image); err != nil {
		return fmt.Errorf("error adding mounts volumes: %s", err)
	}

	return nil
}

func (s *BaseStage) AfterImageSyncDockerStateHook(_ Conveyor) error {
	return nil
}

func (s *BaseStage) PreRunHook(_ Conveyor) error {
	return nil
}

func (s *BaseStage) getServiceMounts(prevBuiltImage image.Image) map[string][]string {
	return mergeMounts(s.getServiceMountsFromLabels(prevBuiltImage), s.getServiceMountsFromConfig())
}

func (s *BaseStage) getServiceMountsFromLabels(prevBuiltImage image.Image) map[string][]string {
	mountpointsByType := map[string][]string{}

	var labels map[string]string
	if prevBuiltImage != nil {
		labels = prevBuiltImage.Labels()
	}

	for _, labelMountType := range []struct{ Label, MountType string }{
		{mountTmpDirLabel, "tmp_dir"},
		{mountBuildDIrLabel, "build_dir"},
	} {
		v, hasKey := labels[labelMountType.Label]
		if !hasKey {
			continue
		}

		mountpoints := util.RejectEmptyStrings(util.UniqStrings(strings.Split(v, ";")))
		mountpointsByType[labelMountType.MountType] = mountpoints
	}

	return mountpointsByType
}

func (s *BaseStage) getServiceMountsFromConfig() map[string][]string {
	mountpointsByType := map[string][]string{}

	for _, mountCfg := range s.configMounts {
		if !util.IsStringsContainValue([]string{"tmp_dir", "build_dir"}, mountCfg.Type) {
			continue
		}

		mountpoint := filepath.Clean(mountCfg.To)
		mountpointsByType[mountCfg.Type] = append(mountpointsByType[mountCfg.Type], mountpoint)
	}

	return mountpointsByType
}

func (s *BaseStage) addServiceMountsVolumes(mountpointsByType map[string][]string, image image.Image) error {
	for mountType, mountpoints := range mountpointsByType {
		for _, mountpoint := range mountpoints {
			absoluteMountpoint := filepath.Join("/", mountpoint)

			var absoluteFrom string
			switch mountType {
			case "tmp_dir":
				absoluteFrom = filepath.Join(s.dimgTmpDir, "mount", slug.Slug(absoluteMountpoint))
			case "build_dir":
				absoluteFrom = filepath.Join(s.projectBuildDir, "mount", slug.Slug(absoluteMountpoint))
			default:
				panic(fmt.Sprintf("unknown mount type %s", mountType))
			}

			err := os.MkdirAll(absoluteFrom, os.ModePerm)
			if err != nil {
				return fmt.Errorf("error creating tmp path %s for mount: %s", absoluteFrom, err)
			}

			image.Container().RunOptions().AddVolume(fmt.Sprintf("%s:%s", absoluteFrom, absoluteMountpoint))
		}
	}

	return nil
}

func (s *BaseStage) addServiceMountsLabels(mountpointsByType map[string][]string, image image.Image) {
	for mountType, mountpoints := range mountpointsByType {
		var labelName string
		switch mountType {
		case "tmp_dir":
			labelName = mountTmpDirLabel
		case "build_dir":
			labelName = mountBuildDIrLabel
		default:
			panic(fmt.Sprintf("unknown mount type %s", mountType))
		}

		labelValue := strings.Join(mountpoints, ";")

		image.Container().ServiceCommitChangeOptions().AddLabel(map[string]string{labelName: labelValue})
	}
}

func (s *BaseStage) getCustomMounts(prevBuiltImage image.Image) map[string][]string {
	return mergeMounts(s.getCustomMountsFromLabels(prevBuiltImage), s.getCustomMountsFromConfig())
}

func (s *BaseStage) getCustomMountsFromLabels(prevBuiltImage image.Image) map[string][]string {
	mountpointsByFrom := map[string][]string{}

	var labels map[string]string
	if prevBuiltImage != nil {
		labels = prevBuiltImage.Labels()
	}
	for k, v := range labels {
		if !strings.HasPrefix(k, mountCustomDirLabelPrefix) {
			continue
		}

		parts := strings.SplitN(k, mountCustomDirLabelPrefix, 2)
		from := strings.Replace(parts[1], "--", "/", -1)

		mountpoints := util.RejectEmptyStrings(util.UniqStrings(strings.Split(v, ";")))
		mountpointsByFrom[from] = mountpoints
	}

	return mountpointsByFrom
}

func (s *BaseStage) getCustomMountsFromConfig() map[string][]string {
	mountpointsByFrom := map[string][]string{}
	for _, mountCfg := range s.configMounts {
		if mountCfg.Type != "custom_dir" {
			continue
		}

		from := filepath.Clean(mountCfg.From)
		mountpoint := filepath.Clean(mountCfg.To)

		mountpointsByFrom[from] = util.UniqAppendString(mountpointsByFrom[from], mountpoint)
	}

	return mountpointsByFrom
}

func (s *BaseStage) addCustomMountVolumes(mountpointsByFrom map[string][]string, image image.Image) error {
	for from, mountpoints := range mountpointsByFrom {
		absoluteFrom := util.ExpandPath(from)

		err := os.MkdirAll(absoluteFrom, os.ModePerm)
		if err != nil {
			return fmt.Errorf("error creating %s: %s", absoluteFrom, err)
		}

		for _, mountpoint := range mountpoints {
			absoluteMountpoint := filepath.Join("/", mountpoint)
			image.Container().RunOptions().AddVolume(fmt.Sprintf("%s:%s", absoluteFrom, absoluteMountpoint))
		}
	}

	return nil
}

func (s *BaseStage) addCustomMountLabels(mountpointsByFrom map[string][]string, image image.Image) {
	for from, mountpoints := range mountpointsByFrom {
		labelName := fmt.Sprintf("%s%s", mountCustomDirLabelPrefix, strings.Replace(from, "/", "--", -1))
		labelValue := strings.Join(mountpoints, ";")
		image.Container().ServiceCommitChangeOptions().AddLabel(map[string]string{labelName: labelValue})
	}
}

func (s *BaseStage) SetSignature(signature string) {
	s.signature = signature
}

func (s *BaseStage) GetSignature() string {
	return s.signature
}

func (s *BaseStage) SetImage(image image.Image) {
	s.image = image
}

func (s *BaseStage) GetImage() image.Image {
	return s.image
}

func (s *BaseStage) SetGitArtifacts(gitArtifacts []*GitArtifact) {
	s.gitArtifacts = gitArtifacts
}

func (s *BaseStage) GetGitArtifacts() []*GitArtifact {
	return s.gitArtifacts
}

func mergeMounts(a, b map[string][]string) map[string][]string {
	res := map[string][]string{}

	for k, mountpoints := range a {
		res[k] = mountpoints
	}
	for k, mountpoints := range b {
		res[k] = util.UniqStrings(append(res[k], mountpoints...))
	}

	return res
}
<commit_msg>Rename artifact import stages<commit_after>package stage

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/flant/dapp/pkg/config"
	"github.com/flant/dapp/pkg/image"
	"github.com/flant/dapp/pkg/slug"
	"github.com/flant/dapp/pkg/util"
)

type StageName string

const (
	From                        StageName = "from"
	BeforeInstall               StageName = "before_install"
	ArtifactImportBeforeInstall StageName = "imports_before_install"
	GAArchive                   StageName = "g_a_archive"
	Install                     StageName = "install"
	ArtifactImportAfterInstall  StageName = "imports_after_install"
	BeforeSetup                 StageName = "before_setup"
	ArtifactImportBeforeSetup   StageName = "imports_before_setup"
	Setup                       StageName = "setup"
	ArtifactImportAfterSetup    StageName = "imports_after_setup"
	GAPostSetupPatch            StageName = "g_a_post_setup_patch"
	GALatestPatch               StageName = "g_a_latest_patch"
	DockerInstructions          StageName = "docker_instructions"
)

const (
	mountTmpDirLabel          = "dapp-mount-type-tmp-dir"
	mountBuildDIrLabel        = "dapp-mount-type-build-dir"
	mountCustomDirLabelPrefix = "dapp-mount-type-custom-dir-"
)

type NewBaseStageOptions struct {
	DimgName         string
	ConfigMounts     []*config.Mount
	DimgTmpDir       string
	ContainerDappDir string
	ProjectBuildDir  string
}

func newBaseStage(name StageName, options *NewBaseStageOptions) *BaseStage {
	s := &BaseStage{}
	s.name = name
	s.dimgName = options.DimgName
	s.configMounts = options.ConfigMounts
	s.projectBuildDir = options.ProjectBuildDir
	s.dimgTmpDir = options.DimgTmpDir
	s.containerDappDir = options.ContainerDappDir
	return s
}

type BaseStage struct {
	name             StageName
	dimgName         string
	signature        string
	image            image.Image
	gitArtifacts     []*GitArtifact
	dimgTmpDir       string
	containerDappDir string
	projectBuildDir  string
	configMounts     []*config.Mount
}

func (s *BaseStage) Name() StageName {
	if s.name != "" {
		return s.name
	}

	panic("name must be defined!")
}

func (s *BaseStage) GetDependencies(_ Conveyor, _ image.Image) (string, error) {
	panic("method must be implemented!")
}

func (s *BaseStage) IsEmpty(_ Conveyor, _ image.Image) (bool, error) {
	return false, nil
}

func (s *BaseStage) ShouldBeReset(_ image.Image) (bool, error) {
	return false, nil
}

func (s *BaseStage) PrepareImage(_ Conveyor, prevBuiltImage, image image.Image) error {
	/*
	 * NOTE: BaseStage.PrepareImage does not called in From.PrepareImage.
	 * NOTE: Take into account when adding new base PrepareImage steps.
	 */

	serviceMounts := s.getServiceMounts(prevBuiltImage)
	s.addServiceMountsLabels(serviceMounts, image)
	if err := s.addServiceMountsVolumes(serviceMounts, image); err != nil {
		return fmt.Errorf("error adding mounts volumes: %s", err)
	}

	customMounts := s.getCustomMounts(prevBuiltImage)
	s.addCustomMountLabels(customMounts, image)
	if err := s.addCustomMountVolumes(customMounts, image); err != nil {
		return fmt.Errorf("error adding mounts volumes: %s", err)
	}

	return nil
}

func (s *BaseStage) AfterImageSyncDockerStateHook(_ Conveyor) error {
	return nil
}

func (s *BaseStage) PreRunHook(_ Conveyor) error {
	return nil
}

func (s *BaseStage) getServiceMounts(prevBuiltImage image.Image) map[string][]string {
	return mergeMounts(s.getServiceMountsFromLabels(prevBuiltImage), s.getServiceMountsFromConfig())
}

func (s *BaseStage) getServiceMountsFromLabels(prevBuiltImage image.Image) map[string][]string {
	mountpointsByType := map[string][]string{}

	var labels map[string]string
	if prevBuiltImage != nil {
		labels = prevBuiltImage.Labels()
	}

	for _, labelMountType := range []struct{ Label, MountType string }{
		{mountTmpDirLabel, "tmp_dir"},
		{mountBuildDIrLabel, "build_dir"},
	} {
		v, hasKey := labels[labelMountType.Label]
		if !hasKey {
			continue
		}

		mountpoints := util.RejectEmptyStrings(util.UniqStrings(strings.Split(v, ";")))
		mountpointsByType[labelMountType.MountType] = mountpoints
	}

	return mountpointsByType
}

func (s *BaseStage) getServiceMountsFromConfig() map[string][]string {
	mountpointsByType := map[string][]string{}

	for _, mountCfg := range s.configMounts {
		if !util.IsStringsContainValue([]string{"tmp_dir", "build_dir"}, mountCfg.Type) {
			continue
		}

		mountpoint := filepath.Clean(mountCfg.To)
		mountpointsByType[mountCfg.Type] = append(mountpointsByType[mountCfg.Type], mountpoint)
	}

	return mountpointsByType
}

func (s *BaseStage) addServiceMountsVolumes(mountpointsByType map[string][]string, image image.Image) error {
	for mountType, mountpoints := range mountpointsByType {
		for _, mountpoint := range mountpoints {
			absoluteMountpoint := filepath.Join("/", mountpoint)

			var absoluteFrom string
			switch mountType {
			case "tmp_dir":
				absoluteFrom = filepath.Join(s.dimgTmpDir, "mount", slug.Slug(absoluteMountpoint))
			case "build_dir":
				absoluteFrom = filepath.Join(s.projectBuildDir, "mount", slug.Slug(absoluteMountpoint))
			default:
				panic(fmt.Sprintf("unknown mount type %s", mountType))
			}

			err := os.MkdirAll(absoluteFrom, os.ModePerm)
			if err != nil {
				return fmt.Errorf("error creating tmp path %s for mount: %s", absoluteFrom, err)
			}

			image.Container().RunOptions().AddVolume(fmt.Sprintf("%s:%s", absoluteFrom, absoluteMountpoint))
		}
	}

	return nil
}

func (s *BaseStage) addServiceMountsLabels(mountpointsByType map[string][]string, image image.Image) {
	for mountType, mountpoints := range mountpointsByType {
		var labelName string
		switch mountType {
		case "tmp_dir":
			labelName = mountTmpDirLabel
		case "build_dir":
			labelName = mountBuildDIrLabel
		default:
			panic(fmt.Sprintf("unknown mount type %s", mountType))
		}

		labelValue := strings.Join(mountpoints, ";")

		image.Container().ServiceCommitChangeOptions().AddLabel(map[string]string{labelName: labelValue})
	}
}

func (s *BaseStage) getCustomMounts(prevBuiltImage image.Image) map[string][]string {
	return mergeMounts(s.getCustomMountsFromLabels(prevBuiltImage), s.getCustomMountsFromConfig())
}

func (s *BaseStage) getCustomMountsFromLabels(prevBuiltImage image.Image) map[string][]string {
	mountpointsByFrom := map[string][]string{}

	var labels map[string]string
	if prevBuiltImage != nil {
		labels = prevBuiltImage.Labels()
	}
	for k, v := range labels {
		if !strings.HasPrefix(k, mountCustomDirLabelPrefix) {
			continue
		}

		parts := strings.SplitN(k, mountCustomDirLabelPrefix, 2)
		from := strings.Replace(parts[1], "--", "/", -1)

		mountpoints := util.RejectEmptyStrings(util.UniqStrings(strings.Split(v, ";")))
		mountpointsByFrom[from] = mountpoints
	}

	return mountpointsByFrom
}

func (s *BaseStage) getCustomMountsFromConfig() map[string][]string {
	mountpointsByFrom := map[string][]string{}
	for _, mountCfg := range s.configMounts {
		if mountCfg.Type != "custom_dir" {
			continue
		}

		from := filepath.Clean(mountCfg.From)
		mountpoint := filepath.Clean(mountCfg.To)

		mountpointsByFrom[from] = util.UniqAppendString(mountpointsByFrom[from], mountpoint)
	}

	return mountpointsByFrom
}

func (s *BaseStage) addCustomMountVolumes(mountpointsByFrom map[string][]string, image image.Image) error {
	for from, mountpoints := range mountpointsByFrom {
		absoluteFrom := util.ExpandPath(from)

		err := os.MkdirAll(absoluteFrom, os.ModePerm)
		if err != nil {
			return fmt.Errorf("error creating %s: %s", absoluteFrom, err)
		}

		for _, mountpoint := range mountpoints {
			absoluteMountpoint := filepath.Join("/", mountpoint)
			image.Container().RunOptions().AddVolume(fmt.Sprintf("%s:%s", absoluteFrom, absoluteMountpoint))
		}
	}

	return nil
}

func (s *BaseStage) addCustomMountLabels(mountpointsByFrom map[string][]string, image image.Image) {
	for from, mountpoints := range mountpointsByFrom {
		labelName := fmt.Sprintf("%s%s", mountCustomDirLabelPrefix, strings.Replace(from, "/", "--", -1))
		labelValue := strings.Join(mountpoints, ";")
		image.Container().ServiceCommitChangeOptions().AddLabel(map[string]string{labelName: labelValue})
	}
}

func (s *BaseStage) SetSignature(signature string) {
	s.signature = signature
}

func (s *BaseStage) GetSignature() string {
	return s.signature
}

func (s *BaseStage) SetImage(image image.Image) {
	s.image = image
}

func (s *BaseStage) GetImage() image.Image {
	return s.image
}

func (s *BaseStage) SetGitArtifacts(gitArtifacts []*GitArtifact) {
	s.gitArtifacts = gitArtifacts
}

func (s *BaseStage) GetGitArtifacts() []*GitArtifact {
	return s.gitArtifacts
}

func mergeMounts(a, b map[string][]string) map[string][]string {
	res := map[string][]string{}

	for k, mountpoints := range a {
		res[k] = mountpoints
	}
	for k, mountpoints := range b {
		res[k] = util.UniqStrings(append(res[k], mountpoints...))
	}

	return res
}
<|endoftext|>


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

// +build !race

package cache

import (
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

type cacheOp int

const (
	Get = iota
	Set
	Remove
	RemoveAll
)

func testCacheBasic(c Cache, t *testing.T) {
	cases := []struct {
		op     cacheOp
		key    string
		value  string
		result bool
		stats  Stats
	}{
		// try to get when the entry isn't present
		{Get, "X", "", false, Stats{Misses: 1}},

		// add an entry and make sure we can get it
		{Set, "X", "12", false, Stats{Misses: 1, Writes: 1}},
		{Get, "X", "12", true, Stats{Misses: 1, Writes: 1, Hits: 1}},
		{Get, "X", "12", true, Stats{Misses: 1, Writes: 1, Hits: 2}},

		// check interference between get/set
		{Get, "Y", "", false, Stats{Misses: 2, Writes: 1, Hits: 2}},
		{Set, "X", "23", false, Stats{Misses: 2, Writes: 2, Hits: 2}},
		{Get, "X", "23", true, Stats{Misses: 2, Writes: 2, Hits: 3}},
		{Set, "Y", "34", false, Stats{Misses: 2, Writes: 3, Hits: 3}},
		{Get, "X", "23", true, Stats{Misses: 2, Writes: 3, Hits: 4}},
		{Get, "Y", "34", true, Stats{Misses: 2, Writes: 3, Hits: 5}},

		// ensure removing X works and doesn't affect Y
		{Remove, "X", "", false, Stats{Misses: 2, Writes: 3, Hits: 5}},
		{Get, "X", "", false, Stats{Misses: 3, Writes: 3, Hits: 5}},
		{Get, "Y", "34", true, Stats{Misses: 3, Writes: 3, Hits: 6}},

		// make sure everything recovers from remove and then get/set
		{Remove, "X", "", false, Stats{Misses: 3, Writes: 3, Hits: 6}},
		{Remove, "Y", "", false, Stats{Misses: 3, Writes: 3, Hits: 6}},
		{Get, "Y", "", false, Stats{Misses: 4, Writes: 3, Hits: 6}},
		{Set, "X", "45", false, Stats{Misses: 4, Writes: 4, Hits: 6}},
		{Get, "X", "45", true, Stats{Misses: 4, Writes: 4, Hits: 7}},
		{Get, "Y", "", false, Stats{Misses: 5, Writes: 4, Hits: 7}},

		// remove a missing entry, should be a nop
		{Remove, "Z", "", false, Stats{Misses: 5, Writes: 4, Hits: 7}},

		// remove everything
		{Set, "A", "45", false, Stats{Misses: 5, Writes: 5, Hits: 7}},
		{Set, "B", "45", false, Stats{Misses: 5, Writes: 6, Hits: 7}},
		{RemoveAll, "", "", false, Stats{Misses: 5, Writes: 6, Hits: 7}},
		{Get, "A", "45", false, Stats{Misses: 6, Writes: 6, Hits: 7}},
		{Get, "B", "45", false, Stats{Misses: 7, Writes: 6, Hits: 7}},
	}

	for i, tc := range cases {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			switch tc.op {
			case Get:
				value, result := c.Get(tc.key)

				if result != tc.result {
					t.Errorf("Got result %v, expected %v", result, tc.result)
				}

				if result {
					str := value.(string)

					if str != tc.value {
						t.Errorf("Got value %v, expected %v", str, tc.value)
					}
				} else {
					if value != nil {
						t.Errorf("Got value %v, expected nil", value)
					}
				}

			case Set:
				c.Set(tc.key, tc.value)

			case Remove:
				c.Remove(tc.key)

			case RemoveAll:
				c.RemoveAll()
			}

			s := c.Stats()

			// removals are inconsistently tracked between implementations, so we ignore these here
			s.Removals = 0

			if s != tc.stats {
				t.Errorf("Got stats of %v, expected %v", s, tc.stats)
			}
		})
	}
}

func testCacheConcurrent(c Cache, t *testing.T) {
	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	wg.Add(workers)

	const numIters = 100000
	for i := 0; i < workers; i++ {
		workerNum := i
		go func() {
			for j := 0; j < numIters; j++ {

				key := "X" + strconv.Itoa(workerNum) + "." + strconv.Itoa(workerNum)
				c.Set(key, j)
				v, ok := c.Get(key)
				if !ok {
					t.Errorf("Got false for key %s, expecting true", key)
				} else if v.(int) != j {
					t.Errorf("Got %d for key %s, expecting %d", v, key, j)
				}
				c.Remove(key)
			}
			wg.Done()
		}()
	}
	wg.Wait()

	stats := c.Stats()
	if stats.Misses != 0 {
		t.Errorf("Got %d misses, expecting %d", stats.Misses, 0)
	}

	if stats.Hits != uint64(workers*numIters) {
		t.Errorf("Got %d hits, expecting %d", stats.Hits, workers*numIters)
	}

	if stats.Writes != uint64(workers*numIters) {
		t.Errorf("Got %d writes, expecting %d", stats.Writes, workers*numIters*2)
	}

}

// WARNING: This test expects the cache to have been created with a long expiration time.
func testCacheExpiration(c ExpiringCache, evictExpired func(time.Time), t *testing.T) {
	now := time.Now()

	c.SetWithExpiration("EARLY", "123", 10*time.Millisecond)
	c.SetWithExpiration("LATER", "123", 20*time.Millisecond+123*time.Nanosecond)

	evictExpired(now)
	s := c.Stats()

	_, ok := c.Get("EARLY")
	if !ok {
		t.Errorf("Got no value, expected EARLY to be present")
	}

	_, ok = c.Get("LATER")
	if !ok {
		t.Errorf("Got no value, expected LATER to be present")
	}

	if s.Evictions != 0 {
		t.Errorf("Got %d evictions, expecting 0", s.Evictions)
	}

	evictExpired(now.Add(15 * time.Millisecond))
	s = c.Stats()

	_, ok = c.Get("EARLY")
	if ok {
		t.Errorf("Got value, expected EARLY to have been evicted")
	}

	_, ok = c.Get("LATER")
	if !ok {
		t.Errorf("Got no value, expected LATER to still be present")
	}

	if s.Evictions != 1 {
		t.Errorf("Got %d evictions, expecting 1", s.Evictions)
	}

	evictExpired(now.Add(25 * time.Millisecond))
	s = c.Stats()

	_, ok = c.Get("EARLY")
	if ok {
		t.Errorf("Got value, expected EARLY to have been evicted")
	}

	_, ok = c.Get("LATER")
	if ok {
		t.Errorf("Got value, expected LATER to have been evicted")
	}

	if s.Evictions != 2 {
		t.Errorf("Got %d evictions, expecting 2", s.Evictions)
	}
}

func testCacheEvictExpired(c ExpiringCache, t *testing.T) {
	c.SetWithExpiration("A", "A", 1*time.Millisecond)

	// this is racy, but we're being generous enough that it should be fine
	time.Sleep(10 * time.Millisecond)

	_, ok := c.Get("A")
	if !ok {
		t.Error("Got no entry, expecting it to be there")
	}

	c.EvictExpired()

	_, ok = c.Get("A")
	if ok {
		t.Error("Got an entry, expecting it to have been evicted")
	}
}

func testCacheEvicter(c ExpiringCache, t *testing.T) {
	c.SetWithExpiration("A", "A", 1*time.Millisecond)

	// this is racy, but we're being generous enough that it should be fine
	time.Sleep(10 * time.Millisecond)

	_, ok := c.Get("A")
	if ok {
		t.Error("Got entry, expecting it to have been evicted")
	}
}

func testCacheFinalizer(gate *bool, t *testing.T) {
	for i := 0; i < 100; i++ {
		runtime.GC()
		if *gate {
			return
		}

		time.Sleep(10 * time.Millisecond)
	}

	t.Errorf("Expecting eviction loop to have been terminated")
}

func benchmarkCacheGet(c Cache, b *testing.B) {
	c.Set("foo", "bar")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.Get("foo")
	}
}

func benchmarkCacheGetConcurrent(c Cache, b *testing.B) {
	c.Set("foo1", "bar")
	c.Set("foo2", "bar")
	c.Set("foo3", "bar")
	c.Set("foo4", "bar")
	c.Set("foo5", "bar")
	c.Set("foo6", "bar")
	c.Set("foo7", "bar")

	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Get("foo1")
				c.Get("foo2")
				c.Get("foo3")
				c.Get("foo5")
				c.Get("foo6")
				c.Get("foo7")
				c.Get("foo8") // doesn't exist
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheSet(c Cache, b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.Set("foo", "bar")
	}
}

func benchmarkCacheSetConcurrent(c Cache, b *testing.B) {
	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Set("foo1", "bar")
				c.Set("foo2", "bar")
				c.Set("foo3", "bar")
				c.Set("foo4", "bar")
				c.Set("foo5", "bar")
				c.Set("foo6", "bar")
				c.Set("foo7", "bar")
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheGetSetConcurrent(c Cache, b *testing.B) {
	c.Set("foo1", "bar")
	c.Set("foo2", "bar")
	c.Set("foo3", "bar")
	c.Set("foo4", "bar")
	c.Set("foo5", "bar")
	c.Set("foo6", "bar")
	c.Set("foo7", "bar")

	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Get("foo1")
				c.Get("foo2")
				c.Get("foo3")
				c.Get("foo5")
				c.Get("foo6")
				c.Get("foo7")
				c.Get("foo8") // doesn't exist

				c.Set("foo1", "bar")
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheSetRemove(c Cache, b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		name := "foo" + strconv.Itoa(i)
		c.Set(name, "bar")
		c.Remove(name)
	}
}
<commit_msg>Bump up the time allowed to from 10ms to 50ms; this is still short, but should save us this test failing ~once a day (#4153)<commit_after>// Copyright 2017 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build !race

package cache

import (
	"runtime"
	"strconv"
	"sync"
	"testing"
	"time"
)

type cacheOp int

const (
	Get = iota
	Set
	Remove
	RemoveAll
)

func testCacheBasic(c Cache, t *testing.T) {
	cases := []struct {
		op     cacheOp
		key    string
		value  string
		result bool
		stats  Stats
	}{
		// try to get when the entry isn't present
		{Get, "X", "", false, Stats{Misses: 1}},

		// add an entry and make sure we can get it
		{Set, "X", "12", false, Stats{Misses: 1, Writes: 1}},
		{Get, "X", "12", true, Stats{Misses: 1, Writes: 1, Hits: 1}},
		{Get, "X", "12", true, Stats{Misses: 1, Writes: 1, Hits: 2}},

		// check interference between get/set
		{Get, "Y", "", false, Stats{Misses: 2, Writes: 1, Hits: 2}},
		{Set, "X", "23", false, Stats{Misses: 2, Writes: 2, Hits: 2}},
		{Get, "X", "23", true, Stats{Misses: 2, Writes: 2, Hits: 3}},
		{Set, "Y", "34", false, Stats{Misses: 2, Writes: 3, Hits: 3}},
		{Get, "X", "23", true, Stats{Misses: 2, Writes: 3, Hits: 4}},
		{Get, "Y", "34", true, Stats{Misses: 2, Writes: 3, Hits: 5}},

		// ensure removing X works and doesn't affect Y
		{Remove, "X", "", false, Stats{Misses: 2, Writes: 3, Hits: 5}},
		{Get, "X", "", false, Stats{Misses: 3, Writes: 3, Hits: 5}},
		{Get, "Y", "34", true, Stats{Misses: 3, Writes: 3, Hits: 6}},

		// make sure everything recovers from remove and then get/set
		{Remove, "X", "", false, Stats{Misses: 3, Writes: 3, Hits: 6}},
		{Remove, "Y", "", false, Stats{Misses: 3, Writes: 3, Hits: 6}},
		{Get, "Y", "", false, Stats{Misses: 4, Writes: 3, Hits: 6}},
		{Set, "X", "45", false, Stats{Misses: 4, Writes: 4, Hits: 6}},
		{Get, "X", "45", true, Stats{Misses: 4, Writes: 4, Hits: 7}},
		{Get, "Y", "", false, Stats{Misses: 5, Writes: 4, Hits: 7}},

		// remove a missing entry, should be a nop
		{Remove, "Z", "", false, Stats{Misses: 5, Writes: 4, Hits: 7}},

		// remove everything
		{Set, "A", "45", false, Stats{Misses: 5, Writes: 5, Hits: 7}},
		{Set, "B", "45", false, Stats{Misses: 5, Writes: 6, Hits: 7}},
		{RemoveAll, "", "", false, Stats{Misses: 5, Writes: 6, Hits: 7}},
		{Get, "A", "45", false, Stats{Misses: 6, Writes: 6, Hits: 7}},
		{Get, "B", "45", false, Stats{Misses: 7, Writes: 6, Hits: 7}},
	}

	for i, tc := range cases {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			switch tc.op {
			case Get:
				value, result := c.Get(tc.key)

				if result != tc.result {
					t.Errorf("Got result %v, expected %v", result, tc.result)
				}

				if result {
					str := value.(string)

					if str != tc.value {
						t.Errorf("Got value %v, expected %v", str, tc.value)
					}
				} else {
					if value != nil {
						t.Errorf("Got value %v, expected nil", value)
					}
				}

			case Set:
				c.Set(tc.key, tc.value)

			case Remove:
				c.Remove(tc.key)

			case RemoveAll:
				c.RemoveAll()
			}

			s := c.Stats()

			// removals are inconsistently tracked between implementations, so we ignore these here
			s.Removals = 0

			if s != tc.stats {
				t.Errorf("Got stats of %v, expected %v", s, tc.stats)
			}
		})
	}
}

func testCacheConcurrent(c Cache, t *testing.T) {
	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	wg.Add(workers)

	const numIters = 100000
	for i := 0; i < workers; i++ {
		workerNum := i
		go func() {
			for j := 0; j < numIters; j++ {

				key := "X" + strconv.Itoa(workerNum) + "." + strconv.Itoa(workerNum)
				c.Set(key, j)
				v, ok := c.Get(key)
				if !ok {
					t.Errorf("Got false for key %s, expecting true", key)
				} else if v.(int) != j {
					t.Errorf("Got %d for key %s, expecting %d", v, key, j)
				}
				c.Remove(key)
			}
			wg.Done()
		}()
	}
	wg.Wait()

	stats := c.Stats()
	if stats.Misses != 0 {
		t.Errorf("Got %d misses, expecting %d", stats.Misses, 0)
	}

	if stats.Hits != uint64(workers*numIters) {
		t.Errorf("Got %d hits, expecting %d", stats.Hits, workers*numIters)
	}

	if stats.Writes != uint64(workers*numIters) {
		t.Errorf("Got %d writes, expecting %d", stats.Writes, workers*numIters*2)
	}

}

// WARNING: This test expects the cache to have been created with a long expiration time.
func testCacheExpiration(c ExpiringCache, evictExpired func(time.Time), t *testing.T) {
	now := time.Now()

	c.SetWithExpiration("EARLY", "123", 10*time.Millisecond)
	c.SetWithExpiration("LATER", "123", 20*time.Millisecond+123*time.Nanosecond)

	evictExpired(now)
	s := c.Stats()

	_, ok := c.Get("EARLY")
	if !ok {
		t.Errorf("Got no value, expected EARLY to be present")
	}

	_, ok = c.Get("LATER")
	if !ok {
		t.Errorf("Got no value, expected LATER to be present")
	}

	if s.Evictions != 0 {
		t.Errorf("Got %d evictions, expecting 0", s.Evictions)
	}

	evictExpired(now.Add(15 * time.Millisecond))
	s = c.Stats()

	_, ok = c.Get("EARLY")
	if ok {
		t.Errorf("Got value, expected EARLY to have been evicted")
	}

	_, ok = c.Get("LATER")
	if !ok {
		t.Errorf("Got no value, expected LATER to still be present")
	}

	if s.Evictions != 1 {
		t.Errorf("Got %d evictions, expecting 1", s.Evictions)
	}

	evictExpired(now.Add(25 * time.Millisecond))
	s = c.Stats()

	_, ok = c.Get("EARLY")
	if ok {
		t.Errorf("Got value, expected EARLY to have been evicted")
	}

	_, ok = c.Get("LATER")
	if ok {
		t.Errorf("Got value, expected LATER to have been evicted")
	}

	if s.Evictions != 2 {
		t.Errorf("Got %d evictions, expecting 2", s.Evictions)
	}
}

func testCacheEvictExpired(c ExpiringCache, t *testing.T) {
	c.SetWithExpiration("A", "A", 1*time.Millisecond)

	// this is racy, but we're being generous enough that it should be fine
	time.Sleep(50 * time.Millisecond)

	_, ok := c.Get("A")
	if !ok {
		t.Error("Got no entry, expecting it to be there")
	}

	c.EvictExpired()

	_, ok = c.Get("A")
	if ok {
		t.Error("Got an entry, expecting it to have been evicted")
	}
}

func testCacheEvicter(c ExpiringCache, t *testing.T) {
	c.SetWithExpiration("A", "A", 1*time.Millisecond)

	// this is racy, but we're being generous enough that it should be fine
	time.Sleep(10 * time.Millisecond)

	_, ok := c.Get("A")
	if ok {
		t.Error("Got entry, expecting it to have been evicted")
	}
}

func testCacheFinalizer(gate *bool, t *testing.T) {
	for i := 0; i < 100; i++ {
		runtime.GC()
		if *gate {
			return
		}

		time.Sleep(10 * time.Millisecond)
	}

	t.Errorf("Expecting eviction loop to have been terminated")
}

func benchmarkCacheGet(c Cache, b *testing.B) {
	c.Set("foo", "bar")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.Get("foo")
	}
}

func benchmarkCacheGetConcurrent(c Cache, b *testing.B) {
	c.Set("foo1", "bar")
	c.Set("foo2", "bar")
	c.Set("foo3", "bar")
	c.Set("foo4", "bar")
	c.Set("foo5", "bar")
	c.Set("foo6", "bar")
	c.Set("foo7", "bar")

	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Get("foo1")
				c.Get("foo2")
				c.Get("foo3")
				c.Get("foo5")
				c.Get("foo6")
				c.Get("foo7")
				c.Get("foo8") // doesn't exist
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheSet(c Cache, b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.Set("foo", "bar")
	}
}

func benchmarkCacheSetConcurrent(c Cache, b *testing.B) {
	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Set("foo1", "bar")
				c.Set("foo2", "bar")
				c.Set("foo3", "bar")
				c.Set("foo4", "bar")
				c.Set("foo5", "bar")
				c.Set("foo6", "bar")
				c.Set("foo7", "bar")
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheGetSetConcurrent(c Cache, b *testing.B) {
	c.Set("foo1", "bar")
	c.Set("foo2", "bar")
	c.Set("foo3", "bar")
	c.Set("foo4", "bar")
	c.Set("foo5", "bar")
	c.Set("foo6", "bar")
	c.Set("foo7", "bar")

	wg := new(sync.WaitGroup)
	workers := runtime.NumCPU()
	each := b.N / workers
	wg.Add(workers)

	b.ResetTimer()
	for i := 0; i < workers; i++ {
		go func() {
			for j := 0; j < each; j++ {
				c.Get("foo1")
				c.Get("foo2")
				c.Get("foo3")
				c.Get("foo5")
				c.Get("foo6")
				c.Get("foo7")
				c.Get("foo8") // doesn't exist

				c.Set("foo1", "bar")
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func benchmarkCacheSetRemove(c Cache, b *testing.B) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		name := "foo" + strconv.Itoa(i)
		c.Set(name, "bar")
		c.Remove(name)
	}
}
<|endoftext|>