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

package workflow_systest

import (
	"errors"
	"log"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
	"github.com/venicegeo/pz-workflow/workflow"
)

type WorkflowTester struct {
	suite.Suite
	client        *workflow.Client
	url           string
	apiKey        string
	uniq          string
	eventTypeId   piazza.Ident
	eventTypeName string
	triggerName   string
	triggerId     piazza.Ident
	serviceId     piazza.Ident
	eventIdY      piazza.Ident
	eventIdN      piazza.Ident
	alertId       piazza.Ident
	jobId         piazza.Ident
	dataId        piazza.Ident
}

var mapType = map[string]interface{}{}
var stringType = "string!"

func (suite *WorkflowTester) setupFixture() {
	t := suite.T()
	assert := assert.New(t)

	var err error

	suite.url = "https://pz-workflow.int.geointservices.io"

	suite.apiKey, err = piazza.GetApiKey("int")
	assert.NoError(err)

	client, err := workflow.NewClient2(suite.url, suite.apiKey)
	assert.NoError(err)
	suite.client = client

	suite.uniq = "systest$" + strconv.Itoa(time.Now().Nanosecond())
	suite.eventTypeName = suite.uniq + "-eventtype"
	suite.triggerName = suite.uniq + "-trigger"
}

func (suite *WorkflowTester) teardownFixture() {
}

func TestRunSuite(t *testing.T) {
	s := &WorkflowTester{}
	suite.Run(t, s)
}

func (suite *WorkflowTester) Test00Init() {
	//t := suite.T()
	//assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()
}

func (suite *WorkflowTester) Test01RegisterService() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	body := map[string]interface{}{
		"url":         "http://pzsvc-hello.int.geointservices.io/hello",
		"contractUrl": "http://pzsvc-hello.int.geointservices.io/contract",
		"method":      "POST",
		"resourceMetadata": map[string]interface{}{
			"name":        "Hello World test",
			"description": "This is the test of Hello World",
			"classType":   "U",
		},
	}

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl: url,
		ApiKey:  suite.apiKey,
		//Preflight:  piazza.SimplePreflight,
		//Postflight: piazza.SimplePostflight,
	}

	obj := map[string]interface{}{}
	code, err := h.Post("/service", body, &obj)
	assert.NoError(err)
	assert.Equal(201, code)
	assert.NotNil(obj)

	assert.IsType(mapType, obj["data"])
	data := obj["data"].(map[string]interface{})
	assert.IsType(stringType, data["serviceId"])
	serviceId := data["serviceId"].(string)
	assert.NotEmpty(serviceId)

	suite.serviceId = piazza.Ident(serviceId)
	log.Printf("ServiceId: %s", suite.serviceId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test02PostEventType() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	eventType := &workflow.EventType{
		Name: suite.eventTypeName,
		Mapping: map[string]elasticsearch.MappingElementTypeName{
			"alpha": elasticsearch.MappingElementTypeString,
			"beta":  elasticsearch.MappingElementTypeInteger,
		},
	}

	ack, err := client.PostEventType(eventType)
	assert.NoError(err)
	assert.NotNil(ack)

	suite.eventTypeId = ack.EventTypeId
	log.Printf("EventTypeId: %s", suite.eventTypeId)

}

func (suite *WorkflowTester) Test03GetEventType() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllEventTypes()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetEventType(suite.eventTypeId)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.eventTypeId, item.EventTypeId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test04PostTrigger() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//suite.eventTypeId = "77bbe4c6-b1ac-4bbb-8e86-0f6e6a731c39"
	//suite.serviceId = "61985d9c-d4d0-45d9-a655-7dcf2dc08fad"

	client := suite.client

	trigger := &workflow.Trigger{
		Name:    suite.triggerName,
		Enabled: true,
		Condition: workflow.Condition{
			EventTypeIds: []piazza.Ident{suite.eventTypeId},
			Query: map[string]interface{}{
				"query": map[string]interface{}{
					"match": map[string]interface{}{
						"beta": 17,
					},
				},
			},
		},
		Job: workflow.JobRequest{
			CreatedBy: "test",
			JobType: workflow.JobType{
				Type: "execute-service",
				Data: map[string]interface{}{
					"dataInputs": map[string]interface{}{
						"": map[string]interface{}{
							"content":  `{"name":"ME", "count":"5"}`,
							"type":     "body",
							"mimeType": "application/json",
						},
					},
					"dataOutput": [](map[string]interface{}){
						{
							"mimeType": "application/json",
							"type":     "text",
						},
					},
					"serviceId": suite.serviceId,
				},
			},
		},
	}

	ack, err := client.PostTrigger(trigger)
	assert.NoError(err)
	assert.NotNil(ack)

	suite.triggerId = ack.TriggerId
	log.Printf("TriggerId: %s", suite.triggerId)
}

func (suite *WorkflowTester) Test05GetTrigger() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllTriggers()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetTrigger(suite.triggerId)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.triggerId, item.TriggerId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test06PostEvent() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	eventY := &workflow.Event{
		EventTypeId: suite.eventTypeId,
		Data: map[string]interface{}{
			"beta":  17,
			"alpha": "quick brown fox",
		},
	}

	eventN := &workflow.Event{
		EventTypeId: suite.eventTypeId,
		Data: map[string]interface{}{
			"beta":  71,
			"alpha": "lazy dog",
		},
	}

	ack, err := client.PostEvent(eventY)
	assert.NoError(err)
	assert.NotNil(ack)
	suite.eventIdY = ack.EventId
	log.Printf("EventIdY: %s", suite.eventIdY)

	ack, err = client.PostEvent(eventN)
	assert.NoError(err)
	assert.NotNil(ack)
	suite.eventIdN = ack.EventId
	log.Printf("EventIdN: %s", suite.eventIdN)
}

func (suite *WorkflowTester) Test07GetEvent() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllEvents()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetEvent(suite.eventIdY)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.eventIdY, item.EventId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test08PostAlert() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	alert := &workflow.Alert{
		TriggerId: "x",
		EventId:   "y",
		JobId:     "z",
	}

	ack, err := client.PostAlert(alert)
	assert.NoError(err)
	assert.NotNil(ack)
}

func (suite *WorkflowTester) Test09GetAlert() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllAlerts()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	items, err = client.GetAlertByTrigger(suite.triggerId)
	assert.NoError(err)
	assert.NotNil(items)
	assert.Len(*items, 1)
	assert.EqualValues(suite.eventIdY, (*items)[0].EventId)

	suite.alertId = (*items)[0].AlertId
	log.Printf("AlertId: %s", suite.alertId)

	suite.jobId = (*items)[0].JobId
	log.Printf("JobId: %s", suite.jobId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test10GetJob() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//client := suite.client

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl:    url,
		ApiKey:     suite.apiKey,
		Preflight:  piazza.SimplePreflight,
		Postflight: piazza.SimplePostflight,
	}

	var data map[string]interface{}

	poll := func() (bool, error) {
		obj := map[string]interface{}{}
		code, err := h.Get("/job/"+string(suite.jobId), &obj)
		if err != nil {
			return false, err
		}
		if code != 200 {
			log.Printf("code is %d", code)
			return false, errors.New("code not 200")
		}
		if obj == nil {
			return false, errors.New("obj was nil")
		}

		var ok bool
		data, ok = obj["data"].(map[string]interface{})
		if !ok {
			return false, errors.New("obj[data] not a map")
		}

		status, ok := data["status"].(string)
		if !ok {
			return false, errors.New("obj[data][status] not a string")
		}

		if status != "Success" {
			return false, nil
		}

		return true, nil
	}

	ok, err := elasticsearch.PollFunction(poll)
	assert.NoError(err)
	assert.True(ok)

	result, ok := data["result"].(map[string]interface{})
	assert.True(ok)
	id, ok := result["dataId"].(string)
	assert.True(ok)

	suite.dataId = piazza.Ident(id)
	log.Printf("DataId: %s", suite.dataId)
}

func (suite *WorkflowTester) Test11GetData() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//client := suite.client

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl: url,
		ApiKey:  suite.apiKey,
		//Preflight:  piazza.SimplePreflight,
		//Postflight: piazza.SimplePostflight,
	}

	obj := map[string]interface{}{}
	code, err := h.Get("/data/"+string(suite.dataId), &obj)
	assert.NoError(err)
	assert.Equal(200, code)
	assert.NotNil(obj)

	var ok bool
	data, ok := obj["data"].(map[string]interface{})
	assert.True(ok)

	log.Printf("## %#v", obj)
	result, ok := data["result"].(map[string]interface{})
	assert.True(ok)
	id, ok := result["dataId"].(string)
	assert.True(ok)
	suite.dataId = piazza.Ident(id)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test99Admin() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	stats, err := client.GetStats()
	assert.NoError(err)

	assert.NotZero(stats.NumEventTypes)
	assert.NotZero(stats.NumEvents)
	assert.NotZero(stats.NumTriggers)
	assert.NotZero(stats.NumAlerts)
	assert.NotZero(stats.NumTriggeredJobs)
}
<commit_msg>works!<commit_after>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package workflow_systest

import (
	"errors"
	"log"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
	"github.com/venicegeo/pz-workflow/workflow"
)

type WorkflowTester struct {
	suite.Suite
	client        *workflow.Client
	url           string
	apiKey        string
	uniq          string
	eventTypeId   piazza.Ident
	eventTypeName string
	triggerName   string
	triggerId     piazza.Ident
	serviceId     piazza.Ident
	eventIdY      piazza.Ident
	eventIdN      piazza.Ident
	alertId       piazza.Ident
	jobId         piazza.Ident
	dataId        piazza.Ident
}

var mapType = map[string]interface{}{}
var stringType = "string!"

func (suite *WorkflowTester) setupFixture() {
	t := suite.T()
	assert := assert.New(t)

	var err error

	suite.url = "https://pz-workflow.int.geointservices.io"

	suite.apiKey, err = piazza.GetApiKey("int")
	assert.NoError(err)

	client, err := workflow.NewClient2(suite.url, suite.apiKey)
	assert.NoError(err)
	suite.client = client

	suite.uniq = "systest$" + strconv.Itoa(time.Now().Nanosecond())
	suite.eventTypeName = suite.uniq + "-eventtype"
	suite.triggerName = suite.uniq + "-trigger"
}

func (suite *WorkflowTester) teardownFixture() {
}

func TestRunSuite(t *testing.T) {
	s := &WorkflowTester{}
	suite.Run(t, s)
}

func (suite *WorkflowTester) Test00Init() {
	//t := suite.T()
	//assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()
}

func (suite *WorkflowTester) Test01RegisterService() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	body := map[string]interface{}{
		"url":         "http://pzsvc-hello.int.geointservices.io/hello",
		"contractUrl": "http://pzsvc-hello.int.geointservices.io/contract",
		"method":      "POST",
		"resourceMetadata": map[string]interface{}{
			"name":        "Hello World test",
			"description": "This is the test of Hello World",
			"classType":   "U",
		},
	}

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl: url,
		ApiKey:  suite.apiKey,
		//Preflight:  piazza.SimplePreflight,
		//Postflight: piazza.SimplePostflight,
	}

	obj := map[string]interface{}{}
	code, err := h.Post("/service", body, &obj)
	assert.NoError(err)
	assert.Equal(201, code)
	assert.NotNil(obj)

	assert.IsType(mapType, obj["data"])
	data := obj["data"].(map[string]interface{})
	assert.IsType(stringType, data["serviceId"])
	serviceId := data["serviceId"].(string)
	assert.NotEmpty(serviceId)

	suite.serviceId = piazza.Ident(serviceId)
	log.Printf("ServiceId: %s", suite.serviceId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test02PostEventType() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	eventType := &workflow.EventType{
		Name: suite.eventTypeName,
		Mapping: map[string]elasticsearch.MappingElementTypeName{
			"alpha": elasticsearch.MappingElementTypeString,
			"beta":  elasticsearch.MappingElementTypeInteger,
		},
	}

	ack, err := client.PostEventType(eventType)
	assert.NoError(err)
	assert.NotNil(ack)

	suite.eventTypeId = ack.EventTypeId
	log.Printf("EventTypeId: %s", suite.eventTypeId)

}

func (suite *WorkflowTester) Test03GetEventType() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllEventTypes()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetEventType(suite.eventTypeId)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.eventTypeId, item.EventTypeId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test04PostTrigger() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//suite.eventTypeId = "77bbe4c6-b1ac-4bbb-8e86-0f6e6a731c39"
	//suite.serviceId = "61985d9c-d4d0-45d9-a655-7dcf2dc08fad"

	client := suite.client

	trigger := &workflow.Trigger{
		Name:    suite.triggerName,
		Enabled: true,
		Condition: workflow.Condition{
			EventTypeIds: []piazza.Ident{suite.eventTypeId},
			Query: map[string]interface{}{
				"query": map[string]interface{}{
					"match": map[string]interface{}{
						"beta": 17,
					},
				},
			},
		},
		Job: workflow.JobRequest{
			CreatedBy: "test",
			JobType: workflow.JobType{
				Type: "execute-service",
				Data: map[string]interface{}{
					"dataInputs": map[string]interface{}{
						"": map[string]interface{}{
							"content":  `{"name":"ME", "count":5}`,
							"type":     "body",
							"mimeType": "application/json",
						},
					},
					"dataOutput": [](map[string]interface{}){
						{
							"mimeType": "application/json",
							"type":     "text",
						},
					},
					"serviceId": suite.serviceId,
				},
			},
		},
	}

	ack, err := client.PostTrigger(trigger)
	assert.NoError(err)
	assert.NotNil(ack)

	suite.triggerId = ack.TriggerId
	log.Printf("TriggerId: %s", suite.triggerId)
}

func (suite *WorkflowTester) Test05GetTrigger() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllTriggers()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetTrigger(suite.triggerId)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.triggerId, item.TriggerId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test06PostEvent() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	eventY := &workflow.Event{
		EventTypeId: suite.eventTypeId,
		Data: map[string]interface{}{
			"beta":  17,
			"alpha": "quick brown fox",
		},
	}

	eventN := &workflow.Event{
		EventTypeId: suite.eventTypeId,
		Data: map[string]interface{}{
			"beta":  71,
			"alpha": "lazy dog",
		},
	}

	ack, err := client.PostEvent(eventY)
	assert.NoError(err)
	assert.NotNil(ack)
	suite.eventIdY = ack.EventId
	log.Printf("EventIdY: %s", suite.eventIdY)

	ack, err = client.PostEvent(eventN)
	assert.NoError(err)
	assert.NotNil(ack)
	suite.eventIdN = ack.EventId
	log.Printf("EventIdN: %s", suite.eventIdN)
}

func (suite *WorkflowTester) Test07GetEvent() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllEvents()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	item, err := client.GetEvent(suite.eventIdY)
	assert.NoError(err)
	assert.NotNil(item)
	assert.EqualValues(suite.eventIdY, item.EventId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test08PostAlert() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	alert := &workflow.Alert{
		TriggerId: "x",
		EventId:   "y",
		JobId:     "z",
	}

	ack, err := client.PostAlert(alert)
	assert.NoError(err)
	assert.NotNil(ack)
}

func (suite *WorkflowTester) Test09GetAlert() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	items, err := client.GetAllAlerts()
	assert.NoError(err)
	assert.True(len(*items) > 1)

	items, err = client.GetAlertByTrigger(suite.triggerId)
	assert.NoError(err)
	assert.NotNil(items)
	assert.Len(*items, 1)
	assert.EqualValues(suite.eventIdY, (*items)[0].EventId)

	suite.alertId = (*items)[0].AlertId
	log.Printf("AlertId: %s", suite.alertId)

	suite.jobId = (*items)[0].JobId
	log.Printf("JobId: %s", suite.jobId)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test10GetJob() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//client := suite.client

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl: url,
		ApiKey:  suite.apiKey,
		//Preflight:  piazza.SimplePreflight,
		//Postflight: piazza.SimplePostflight,
	}

	var data map[string]interface{}

	poll := func() (bool, error) {
		obj := map[string]interface{}{}
		code, err := h.Get("/job/"+string(suite.jobId), &obj)
		if err != nil {
			return false, err
		}
		if code != 200 {
			log.Printf("code is %d", code)
			return false, errors.New("code not 200")
		}
		if obj == nil {
			return false, errors.New("obj was nil")
		}

		var ok bool
		data, ok = obj["data"].(map[string]interface{})
		if !ok {
			return false, errors.New("obj[data] not a map")
		}

		status, ok := data["status"].(string)
		if !ok {
			return false, errors.New("obj[data][status] not a string")
		}

		if status != "Success" {
			return false, nil
		}

		return true, nil
	}

	ok, err := elasticsearch.PollFunction(poll)
	assert.NoError(err)
	assert.True(ok)

	result, ok := data["result"].(map[string]interface{})
	assert.True(ok)
	id, ok := result["dataId"].(string)
	assert.True(ok)

	suite.dataId = piazza.Ident(id)
	log.Printf("DataId: %s", suite.dataId)
}

func (suite *WorkflowTester) Test11GetData() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	//client := suite.client

	url := strings.Replace(suite.url, "workflow", "gateway", 1)
	h := piazza.Http{
		BaseUrl: url,
		ApiKey:  suite.apiKey,
		//Preflight:  piazza.SimplePreflight,
		//Postflight: piazza.SimplePostflight,
	}

	obj := map[string]interface{}{}
	code, err := h.Get("/data/"+string(suite.dataId), &obj)
	assert.NoError(err)
	assert.Equal(200, code)
	assert.NotNil(obj)

	var ok bool
	data, ok := obj["data"].(map[string]interface{})
	assert.True(ok)

	dataType, ok := data["dataType"].(map[string]interface{})
	assert.True(ok)
	content, ok := dataType["content"].(string)
	assert.True(ok)

	jsn := `{
		"greeting": "Hello, ME!",
		"countSquared": 25
	}`
	assert.JSONEq(jsn, content)
}

//---------------------------------------------------------------------

func (suite *WorkflowTester) Test99Admin() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	stats, err := client.GetStats()
	assert.NoError(err)

	assert.NotZero(stats.NumEventTypes)
	assert.NotZero(stats.NumEvents)
	assert.NotZero(stats.NumTriggers)
	assert.NotZero(stats.NumAlerts)
	assert.NotZero(stats.NumTriggeredJobs)
}
<|endoftext|>


<commit_before>// +build linux

/*
   Copyright The containerd Authors.

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

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

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

package devmapper

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/containerd/containerd/pkg/testutil"
	"github.com/containerd/containerd/snapshots/devmapper/dmsetup"
	"github.com/containerd/containerd/snapshots/devmapper/losetup"
	"github.com/docker/go-units"
	"github.com/sirupsen/logrus"
	"gotest.tools/assert"
)

const (
	thinDevice1 = "thin-1"
	thinDevice2 = "thin-2"
	snapDevice1 = "snap-1"
	device1Size = 100000
	device2Size = 200000
	testsPrefix = "devmapper-snapshotter-tests-"
)

// TestPoolDevice runs integration tests for pool device.
// The following scenario implemented:
// - Create pool device with name 'test-pool-device'
// - Create two thin volumes 'thin-1' and 'thin-2'
// - Write ext4 file system on 'thin-1' and make sure it'errs moutable
// - Write v1 test file on 'thin-1' volume
// - Take 'thin-1' snapshot 'snap-1'
// - Change v1 file to v2 on 'thin-1'
// - Mount 'snap-1' and make sure test file is v1
// - Unmount volumes and remove all devices
func TestPoolDevice(t *testing.T) {
	testutil.RequiresRoot(t)

	logrus.SetLevel(logrus.DebugLevel)
	ctx := context.Background()

	tempDir, err := ioutil.TempDir("", "pool-device-test-")
	assert.NilError(t, err, "couldn't get temp directory for testing")

	_, loopDataDevice := createLoopbackDevice(t, tempDir)
	_, loopMetaDevice := createLoopbackDevice(t, tempDir)

	poolName := fmt.Sprintf("test-pool-device-%d", time.Now().Nanosecond())
	err = dmsetup.CreatePool(poolName, loopDataDevice, loopMetaDevice, 64*1024/dmsetup.SectorSize)
	assert.NilError(t, err, "failed to create pool %q", poolName)

	defer func() {
		// Detach loop devices and remove images
		err := losetup.DetachLoopDevice(loopDataDevice, loopMetaDevice)
		assert.NilError(t, err)

		err = os.RemoveAll(tempDir)
		assert.NilError(t, err, "couldn't cleanup temp directory")
	}()

	config := &Config{
		PoolName:           poolName,
		RootPath:           tempDir,
		BaseImageSize:      "16mb",
		BaseImageSizeBytes: 16 * 1024 * 1024,
	}

	pool, err := NewPoolDevice(ctx, config)
	assert.NilError(t, err, "can't create device pool")
	assert.Assert(t, pool != nil)

	defer func() {
		err := pool.RemovePool(ctx)
		assert.NilError(t, err, "can't close device pool")
	}()

	// Create thin devices
	t.Run("CreateThinDevice", func(t *testing.T) {
		testCreateThinDevice(t, pool)
	})

	// Make ext4 filesystem on 'thin-1'
	t.Run("MakeFileSystem", func(t *testing.T) {
		testMakeFileSystem(t, pool)
	})

	// Mount 'thin-1'
	thin1MountPath := tempMountPath(t)
	output, err := exec.Command("mount", dmsetup.GetFullDevicePath(thinDevice1), thin1MountPath).CombinedOutput()
	assert.NilError(t, err, "failed to mount '%s': %s", thinDevice1, string(output))

	// Write v1 test file on 'thin-1' device
	thin1TestFilePath := filepath.Join(thin1MountPath, "TEST")
	err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v1)"), 0700)
	assert.NilError(t, err, "failed to write test file v1 on '%s' volume", thinDevice1)

	// Take snapshot of 'thin-1'
	t.Run("CreateSnapshotDevice", func(t *testing.T) {
		testCreateSnapshot(t, pool)
	})

	// Update TEST file on 'thin-1' to v2
	err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v2)"), 0700)
	assert.NilError(t, err, "failed to write test file v2 on 'thin-1' volume after taking snapshot")

	// Mount 'snap-1' and make sure TEST file is v1
	snap1MountPath := tempMountPath(t)
	output, err = exec.Command("mount", dmsetup.GetFullDevicePath(snapDevice1), snap1MountPath).CombinedOutput()
	assert.NilError(t, err, "failed to mount '%s' device: %s", snapDevice1, string(output))

	// Read test file from snapshot device and make sure it's v1
	fileData, err := ioutil.ReadFile(filepath.Join(snap1MountPath, "TEST"))
	assert.NilError(t, err, "couldn't read test file from '%s' device", snapDevice1)
	assert.Assert(t, string(fileData) == "test file (v1)", "test file content is invalid on snapshot")

	// Unmount devices before removing
	output, err = exec.Command("umount", thin1MountPath, snap1MountPath).CombinedOutput()
	assert.NilError(t, err, "failed to unmount devices: %s", string(output))

	t.Run("DeactivateDevice", func(t *testing.T) {
		testDeactivateThinDevice(t, pool)
	})

	t.Run("RemoveDevice", func(t *testing.T) {
		testRemoveThinDevice(t, pool)
	})
}

func testCreateThinDevice(t *testing.T, pool *PoolDevice) {
	ctx := context.Background()

	err := pool.CreateThinDevice(ctx, thinDevice1, device1Size)
	assert.NilError(t, err, "can't create first thin device")

	err = pool.CreateThinDevice(ctx, thinDevice1, device1Size)
	assert.Assert(t, err != nil, "device pool allows duplicated device names")

	err = pool.CreateThinDevice(ctx, thinDevice2, device2Size)
	assert.NilError(t, err, "can't create second thin device")

	deviceInfo1, err := pool.metadata.GetDevice(ctx, thinDevice1)
	assert.NilError(t, err)

	deviceInfo2, err := pool.metadata.GetDevice(ctx, thinDevice2)
	assert.NilError(t, err)

	assert.Assert(t, deviceInfo1.DeviceID != deviceInfo2.DeviceID, "assigned device ids should be different")
}

func testMakeFileSystem(t *testing.T, pool *PoolDevice) {
	devicePath := dmsetup.GetFullDevicePath(thinDevice1)
	args := []string{
		devicePath,
		"-E",
		"nodiscard,lazy_itable_init=0,lazy_journal_init=0",
	}

	output, err := exec.Command("mkfs.ext4", args...).CombinedOutput()
	assert.NilError(t, err, "failed to make filesystem on '%s': %s", thinDevice1, string(output))
}

func testCreateSnapshot(t *testing.T, pool *PoolDevice) {
	err := pool.CreateSnapshotDevice(context.Background(), thinDevice1, snapDevice1, device1Size)
	assert.NilError(t, err, "failed to create snapshot from '%s' volume", thinDevice1)
}

func testDeactivateThinDevice(t *testing.T, pool *PoolDevice) {
	deviceList := []string{
		thinDevice2,
		snapDevice1,
	}

	for _, deviceName := range deviceList {
		assert.Assert(t, pool.IsActivated(deviceName))

		err := pool.DeactivateDevice(context.Background(), deviceName, false)
		assert.NilError(t, err, "failed to remove '%s'", deviceName)

		assert.Assert(t, !pool.IsActivated(deviceName))
	}
}

func testRemoveThinDevice(t *testing.T, pool *PoolDevice) {
	err := pool.RemoveDevice(testCtx, thinDevice1)
	assert.NilError(t, err, "should delete thin device from pool")
}

func tempMountPath(t *testing.T) string {
	path, err := ioutil.TempDir("", "devmapper-snapshotter-mount-")
	assert.NilError(t, err, "failed to get temp directory for mount")

	return path
}

func createLoopbackDevice(t *testing.T, dir string) (string, string) {
	file, err := ioutil.TempFile(dir, testsPrefix)
	assert.NilError(t, err)

	size, err := units.RAMInBytes("128Mb")
	assert.NilError(t, err)

	err = file.Truncate(size)
	assert.NilError(t, err)

	err = file.Close()
	assert.NilError(t, err)

	imagePath := file.Name()

	loopDevice, err := losetup.AttachLoopDevice(imagePath)
	assert.NilError(t, err)

	return imagePath, loopDevice
}
<commit_msg>devmapper: proper cleanup in pool device test<commit_after>// +build linux

/*
   Copyright The containerd Authors.

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

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

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

package devmapper

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/pkg/testutil"
	"github.com/containerd/containerd/snapshots/devmapper/dmsetup"
	"github.com/containerd/containerd/snapshots/devmapper/losetup"
	"github.com/docker/go-units"
	"github.com/sirupsen/logrus"
	"gotest.tools/assert"
)

const (
	thinDevice1 = "thin-1"
	thinDevice2 = "thin-2"
	snapDevice1 = "snap-1"
	device1Size = 100000
	device2Size = 200000
	testsPrefix = "devmapper-snapshotter-tests-"
)

// TestPoolDevice runs integration tests for pool device.
// The following scenario implemented:
// - Create pool device with name 'test-pool-device'
// - Create two thin volumes 'thin-1' and 'thin-2'
// - Write ext4 file system on 'thin-1' and make sure it'errs moutable
// - Write v1 test file on 'thin-1' volume
// - Take 'thin-1' snapshot 'snap-1'
// - Change v1 file to v2 on 'thin-1'
// - Mount 'snap-1' and make sure test file is v1
// - Unmount volumes and remove all devices
func TestPoolDevice(t *testing.T) {
	testutil.RequiresRoot(t)

	logrus.SetLevel(logrus.DebugLevel)
	ctx := context.Background()

	tempDir, err := ioutil.TempDir("", "pool-device-test-")
	assert.NilError(t, err, "couldn't get temp directory for testing")

	_, loopDataDevice := createLoopbackDevice(t, tempDir)
	_, loopMetaDevice := createLoopbackDevice(t, tempDir)

	poolName := fmt.Sprintf("test-pool-device-%d", time.Now().Nanosecond())
	err = dmsetup.CreatePool(poolName, loopDataDevice, loopMetaDevice, 64*1024/dmsetup.SectorSize)
	assert.NilError(t, err, "failed to create pool %q", poolName)

	defer func() {
		// Detach loop devices and remove images
		err := losetup.DetachLoopDevice(loopDataDevice, loopMetaDevice)
		assert.NilError(t, err)

		err = os.RemoveAll(tempDir)
		assert.NilError(t, err, "couldn't cleanup temp directory")
	}()

	config := &Config{
		PoolName:           poolName,
		RootPath:           tempDir,
		BaseImageSize:      "16mb",
		BaseImageSizeBytes: 16 * 1024 * 1024,
	}

	pool, err := NewPoolDevice(ctx, config)
	assert.NilError(t, err, "can't create device pool")
	assert.Assert(t, pool != nil)

	defer func() {
		err := pool.RemovePool(ctx)
		assert.NilError(t, err, "can't close device pool")
	}()

	// Create thin devices
	t.Run("CreateThinDevice", func(t *testing.T) {
		testCreateThinDevice(t, pool)
	})

	// Make ext4 filesystem on 'thin-1'
	t.Run("MakeFileSystem", func(t *testing.T) {
		testMakeFileSystem(t, pool)
	})

	// Mount 'thin-1'
	err = mount.WithTempMount(ctx, getMounts(thinDevice1), func(thin1MountPath string) error {
		// Write v1 test file on 'thin-1' device
		thin1TestFilePath := filepath.Join(thin1MountPath, "TEST")
		err := ioutil.WriteFile(thin1TestFilePath, []byte("test file (v1)"), 0700)
		assert.NilError(t, err, "failed to write test file v1 on '%s' volume", thinDevice1)

		// Take snapshot of 'thin-1'
		t.Run("CreateSnapshotDevice", func(t *testing.T) {
			testCreateSnapshot(t, pool)
		})

		// Update TEST file on 'thin-1' to v2
		err = ioutil.WriteFile(thin1TestFilePath, []byte("test file (v2)"), 0700)
		assert.NilError(t, err, "failed to write test file v2 on 'thin-1' volume after taking snapshot")

		return nil
	})

	assert.NilError(t, err)

	// Mount 'snap-1' and make sure TEST file is v1
	err = mount.WithTempMount(ctx, getMounts(snapDevice1), func(snap1MountPath string) error {
		// Read test file from snapshot device and make sure it's v1
		fileData, err := ioutil.ReadFile(filepath.Join(snap1MountPath, "TEST"))
		assert.NilError(t, err, "couldn't read test file from '%s' device", snapDevice1)
		assert.Equal(t, "test file (v1)", string(fileData), "test file content is invalid on snapshot")

		return nil
	})

	assert.NilError(t, err)

	t.Run("DeactivateDevice", func(t *testing.T) {
		testDeactivateThinDevice(t, pool)
	})

	t.Run("RemoveDevice", func(t *testing.T) {
		testRemoveThinDevice(t, pool)
	})
}

func testCreateThinDevice(t *testing.T, pool *PoolDevice) {
	ctx := context.Background()

	err := pool.CreateThinDevice(ctx, thinDevice1, device1Size)
	assert.NilError(t, err, "can't create first thin device")

	err = pool.CreateThinDevice(ctx, thinDevice1, device1Size)
	assert.Assert(t, err != nil, "device pool allows duplicated device names")

	err = pool.CreateThinDevice(ctx, thinDevice2, device2Size)
	assert.NilError(t, err, "can't create second thin device")

	deviceInfo1, err := pool.metadata.GetDevice(ctx, thinDevice1)
	assert.NilError(t, err)

	deviceInfo2, err := pool.metadata.GetDevice(ctx, thinDevice2)
	assert.NilError(t, err)

	assert.Assert(t, deviceInfo1.DeviceID != deviceInfo2.DeviceID, "assigned device ids should be different")
}

func testMakeFileSystem(t *testing.T, pool *PoolDevice) {
	devicePath := dmsetup.GetFullDevicePath(thinDevice1)
	args := []string{
		devicePath,
		"-E",
		"nodiscard,lazy_itable_init=0,lazy_journal_init=0",
	}

	output, err := exec.Command("mkfs.ext4", args...).CombinedOutput()
	assert.NilError(t, err, "failed to make filesystem on '%s': %s", thinDevice1, string(output))
}

func testCreateSnapshot(t *testing.T, pool *PoolDevice) {
	err := pool.CreateSnapshotDevice(context.Background(), thinDevice1, snapDevice1, device1Size)
	assert.NilError(t, err, "failed to create snapshot from '%s' volume", thinDevice1)
}

func testDeactivateThinDevice(t *testing.T, pool *PoolDevice) {
	deviceList := []string{
		thinDevice2,
		snapDevice1,
	}

	for _, deviceName := range deviceList {
		assert.Assert(t, pool.IsActivated(deviceName))

		err := pool.DeactivateDevice(context.Background(), deviceName, false)
		assert.NilError(t, err, "failed to remove '%s'", deviceName)

		assert.Assert(t, !pool.IsActivated(deviceName))
	}
}

func testRemoveThinDevice(t *testing.T, pool *PoolDevice) {
	err := pool.RemoveDevice(testCtx, thinDevice1)
	assert.NilError(t, err, "should delete thin device from pool")
}

func getMounts(thinDeviceName string) []mount.Mount {
	return []mount.Mount{
		{
			Source: dmsetup.GetFullDevicePath(thinDeviceName),
			Type:   "ext4",
		},
	}
}

func createLoopbackDevice(t *testing.T, dir string) (string, string) {
	file, err := ioutil.TempFile(dir, testsPrefix)
	assert.NilError(t, err)

	size, err := units.RAMInBytes("128Mb")
	assert.NilError(t, err)

	err = file.Truncate(size)
	assert.NilError(t, err)

	err = file.Close()
	assert.NilError(t, err)

	imagePath := file.Name()

	loopDevice, err := losetup.AttachLoopDevice(imagePath)
	assert.NilError(t, err)

	return imagePath, loopDevice
}
<|endoftext|>


<commit_before>// +build linux

/*
   Copyright The containerd Authors.

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

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

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

package devmapper

import (
	"context"
	_ "crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	"github.com/containerd/continuity/fs/fstest"
	"github.com/hashicorp/go-multierror"
	"github.com/sirupsen/logrus"
	"gotest.tools/v3/assert"

	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/pkg/testutil"
	"github.com/containerd/containerd/snapshots"
	"github.com/containerd/containerd/snapshots/devmapper/dmsetup"
	"github.com/containerd/containerd/snapshots/testsuite"
)

func TestSnapshotterSuite(t *testing.T) {
	testutil.RequiresRoot(t)

	logrus.SetLevel(logrus.DebugLevel)

	snapshotterFn := func(ctx context.Context, root string) (snapshots.Snapshotter, func() error, error) {
		// Create loopback devices for each test case
		_, loopDataDevice := createLoopbackDevice(t, root)
		_, loopMetaDevice := createLoopbackDevice(t, root)

		poolName := fmt.Sprintf("containerd-snapshotter-suite-pool-%d", time.Now().Nanosecond())
		err := dmsetup.CreatePool(poolName, loopDataDevice, loopMetaDevice, 64*1024/dmsetup.SectorSize)
		assert.NilError(t, err, "failed to create pool %q", poolName)

		config := &Config{
			RootPath:      root,
			PoolName:      poolName,
			BaseImageSize: "16Mb",
		}

		snap, err := NewSnapshotter(context.Background(), config)
		if err != nil {
			return nil, nil, err
		}

		// Remove device mapper pool and detach loop devices after test completes
		removePool := func() error {
			result := multierror.Append(
				snap.pool.RemovePool(ctx),
				mount.DetachLoopDevice(loopDataDevice, loopMetaDevice))

			return result.ErrorOrNil()
		}

		// Pool cleanup should be called before closing metadata store (as we need to retrieve device names)
		snap.cleanupFn = append([]closeFunc{removePool}, snap.cleanupFn...)

		return snap, snap.Close, nil
	}

	testsuite.SnapshotterSuite(t, "devmapper", snapshotterFn)

	ctx := context.Background()
	ctx = namespaces.WithNamespace(ctx, "testsuite")

	t.Run("DevMapperUsage", func(t *testing.T) {
		tempDir, err := ioutil.TempDir("", "snapshot-suite-usage")
		assert.NilError(t, err)
		defer os.RemoveAll(tempDir)

		snapshotter, closer, err := snapshotterFn(ctx, tempDir)
		assert.NilError(t, err)
		defer closer()

		testUsage(t, snapshotter)
	})
}

// testUsage tests devmapper's Usage implementation. This is an approximate test as it's hard to
// predict how many blocks will be consumed under different conditions and parameters.
func testUsage(t *testing.T, snapshotter snapshots.Snapshotter) {
	ctx := context.Background()

	// Create empty base layer
	_, err := snapshotter.Prepare(ctx, "prepare-1", "")
	assert.NilError(t, err)

	emptyLayerUsage, err := snapshotter.Usage(ctx, "prepare-1")
	assert.NilError(t, err)

	// Should be > 0 as just written file system also consumes blocks
	assert.Assert(t, emptyLayerUsage.Size > 0)

	err = snapshotter.Commit(ctx, "layer-1", "prepare-1")
	assert.NilError(t, err)

	// Create child layer with 1MB file

	var (
		sizeBytes   int64 = 1048576 // 1MB
		baseApplier       = fstest.Apply(fstest.CreateRandomFile("/a", 12345679, sizeBytes, 0777))
	)

	mounts, err := snapshotter.Prepare(ctx, "prepare-2", "layer-1")
	assert.NilError(t, err)

	err = mount.WithTempMount(ctx, mounts, baseApplier.Apply)
	assert.NilError(t, err)

	err = snapshotter.Commit(ctx, "layer-2", "prepare-2")
	assert.NilError(t, err)

	layer2Usage, err := snapshotter.Usage(ctx, "layer-2")
	assert.NilError(t, err)

	// Should be at least 1 MB + fs metadata
	assert.Check(t, layer2Usage.Size > sizeBytes,
		"%d > %d", layer2Usage.Size > sizeBytes)
	assert.Check(t, layer2Usage.Size < sizeBytes+256*dmsetup.SectorSize,
		"%d < %d", layer2Usage.Size < sizeBytes+256*dmsetup.SectorSize)
}
<commit_msg>test: fix assert.Check's argumets to show its parameters correctly<commit_after>// +build linux

/*
   Copyright The containerd Authors.

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

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

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

package devmapper

import (
	"context"
	_ "crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	"github.com/containerd/continuity/fs/fstest"
	"github.com/hashicorp/go-multierror"
	"github.com/sirupsen/logrus"
	"gotest.tools/v3/assert"

	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/pkg/testutil"
	"github.com/containerd/containerd/snapshots"
	"github.com/containerd/containerd/snapshots/devmapper/dmsetup"
	"github.com/containerd/containerd/snapshots/testsuite"
)

func TestSnapshotterSuite(t *testing.T) {
	testutil.RequiresRoot(t)

	logrus.SetLevel(logrus.DebugLevel)

	snapshotterFn := func(ctx context.Context, root string) (snapshots.Snapshotter, func() error, error) {
		// Create loopback devices for each test case
		_, loopDataDevice := createLoopbackDevice(t, root)
		_, loopMetaDevice := createLoopbackDevice(t, root)

		poolName := fmt.Sprintf("containerd-snapshotter-suite-pool-%d", time.Now().Nanosecond())
		err := dmsetup.CreatePool(poolName, loopDataDevice, loopMetaDevice, 64*1024/dmsetup.SectorSize)
		assert.NilError(t, err, "failed to create pool %q", poolName)

		config := &Config{
			RootPath:      root,
			PoolName:      poolName,
			BaseImageSize: "16Mb",
		}

		snap, err := NewSnapshotter(context.Background(), config)
		if err != nil {
			return nil, nil, err
		}

		// Remove device mapper pool and detach loop devices after test completes
		removePool := func() error {
			result := multierror.Append(
				snap.pool.RemovePool(ctx),
				mount.DetachLoopDevice(loopDataDevice, loopMetaDevice))

			return result.ErrorOrNil()
		}

		// Pool cleanup should be called before closing metadata store (as we need to retrieve device names)
		snap.cleanupFn = append([]closeFunc{removePool}, snap.cleanupFn...)

		return snap, snap.Close, nil
	}

	testsuite.SnapshotterSuite(t, "devmapper", snapshotterFn)

	ctx := context.Background()
	ctx = namespaces.WithNamespace(ctx, "testsuite")

	t.Run("DevMapperUsage", func(t *testing.T) {
		tempDir, err := ioutil.TempDir("", "snapshot-suite-usage")
		assert.NilError(t, err)
		defer os.RemoveAll(tempDir)

		snapshotter, closer, err := snapshotterFn(ctx, tempDir)
		assert.NilError(t, err)
		defer closer()

		testUsage(t, snapshotter)
	})
}

// testUsage tests devmapper's Usage implementation. This is an approximate test as it's hard to
// predict how many blocks will be consumed under different conditions and parameters.
func testUsage(t *testing.T, snapshotter snapshots.Snapshotter) {
	ctx := context.Background()

	// Create empty base layer
	_, err := snapshotter.Prepare(ctx, "prepare-1", "")
	assert.NilError(t, err)

	emptyLayerUsage, err := snapshotter.Usage(ctx, "prepare-1")
	assert.NilError(t, err)

	// Should be > 0 as just written file system also consumes blocks
	assert.Assert(t, emptyLayerUsage.Size > 0)

	err = snapshotter.Commit(ctx, "layer-1", "prepare-1")
	assert.NilError(t, err)

	// Create child layer with 1MB file

	var (
		sizeBytes   int64 = 1048576 // 1MB
		baseApplier       = fstest.Apply(fstest.CreateRandomFile("/a", 12345679, sizeBytes, 0777))
	)

	mounts, err := snapshotter.Prepare(ctx, "prepare-2", "layer-1")
	assert.NilError(t, err)

	err = mount.WithTempMount(ctx, mounts, baseApplier.Apply)
	assert.NilError(t, err)

	err = snapshotter.Commit(ctx, "layer-2", "prepare-2")
	assert.NilError(t, err)

	layer2Usage, err := snapshotter.Usage(ctx, "layer-2")
	assert.NilError(t, err)

	// Should be at least 1 MB + fs metadata
	assert.Check(t, layer2Usage.Size > sizeBytes,
		"%d > %d", layer2Usage.Size, sizeBytes)
	assert.Check(t, layer2Usage.Size < sizeBytes+256*dmsetup.SectorSize,
		"%d < %d", layer2Usage.Size, sizeBytes+256*dmsetup.SectorSize)
}
<|endoftext|>


<commit_before>// Copyright (C) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package router

import (
	"context"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/message"
	"github.com/ava-labs/avalanchego/snow/networking/handler"
	"github.com/ava-labs/avalanchego/snow/networking/timeout"
	"github.com/ava-labs/avalanchego/trace"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/version"
)

var _ Router = (*tracedRouter)(nil)

type tracedRouter struct {
	router Router
	tracer trace.Tracer
}

func Trace(router Router, tracer trace.Tracer) Router {
	return &tracedRouter{
		router: router,
		tracer: tracer,
	}
}

func (r *tracedRouter) Initialize(
	nodeID ids.NodeID,
	log logging.Logger,
	timeoutManager timeout.Manager,
	closeTimeout time.Duration,
	criticalChains ids.Set,
	whitelistedSubnets ids.Set,
	onFatal func(exitCode int),
	healthConfig HealthConfig,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) error {
	return r.router.Initialize(
		nodeID,
		log,
		timeoutManager,
		closeTimeout,
		criticalChains,
		whitelistedSubnets,
		onFatal,
		healthConfig,
		metricsNamespace,
		metricsRegisterer,
	)
}

func (r *tracedRouter) RegisterRequest(
	ctx context.Context,
	nodeID ids.NodeID,
	requestingChainID ids.ID,
	respondingChainID ids.ID,
	requestID uint32,
	op message.Op,
	failedMsg message.InboundMessage,
) {
	r.router.RegisterRequest(
		ctx,
		nodeID,
		requestingChainID,
		respondingChainID,
		requestID,
		op,
		failedMsg,
	)
}

func (r *tracedRouter) HandleInbound(ctx context.Context, msg message.InboundMessage) {
	ctx, span := r.tracer.Start(ctx, "tracedRouter.HandleInbound")
	defer span.End()

	r.router.HandleInbound(ctx, msg)
}

func (r *tracedRouter) Shutdown() {
	r.router.Shutdown()
}

func (r *tracedRouter) AddChain(chain handler.Handler) {
	r.router.AddChain(chain)
}

func (r *tracedRouter) Connected(nodeID ids.NodeID, nodeVersion *version.Application, subnetID ids.ID) {
	r.router.Connected(nodeID, nodeVersion, subnetID)
}

func (r *tracedRouter) Disconnected(nodeID ids.NodeID) {
	r.router.Disconnected(nodeID)
}

func (r *tracedRouter) Benched(chainID ids.ID, nodeID ids.NodeID) {
	r.router.Benched(chainID, nodeID)
}

func (r *tracedRouter) Unbenched(chainID ids.ID, nodeID ids.NodeID) {
	r.router.Unbenched(chainID, nodeID)
}

func (r *tracedRouter) HealthCheck() (interface{}, error) {
	return r.router.HealthCheck()
}
<commit_msg>Improve router tracing (#2208)<commit_after>// Copyright (C) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package router

import (
	"context"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"go.opentelemetry.io/otel/attribute"

	oteltrace "go.opentelemetry.io/otel/trace"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/message"
	"github.com/ava-labs/avalanchego/snow/networking/handler"
	"github.com/ava-labs/avalanchego/snow/networking/timeout"
	"github.com/ava-labs/avalanchego/trace"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/version"
)

var _ Router = (*tracedRouter)(nil)

type tracedRouter struct {
	router Router
	tracer trace.Tracer
}

func Trace(router Router, tracer trace.Tracer) Router {
	return &tracedRouter{
		router: router,
		tracer: tracer,
	}
}

func (r *tracedRouter) Initialize(
	nodeID ids.NodeID,
	log logging.Logger,
	timeoutManager timeout.Manager,
	closeTimeout time.Duration,
	criticalChains ids.Set,
	whitelistedSubnets ids.Set,
	onFatal func(exitCode int),
	healthConfig HealthConfig,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) error {
	return r.router.Initialize(
		nodeID,
		log,
		timeoutManager,
		closeTimeout,
		criticalChains,
		whitelistedSubnets,
		onFatal,
		healthConfig,
		metricsNamespace,
		metricsRegisterer,
	)
}

func (r *tracedRouter) RegisterRequest(
	ctx context.Context,
	nodeID ids.NodeID,
	requestingChainID ids.ID,
	respondingChainID ids.ID,
	requestID uint32,
	op message.Op,
	failedMsg message.InboundMessage,
) {
	r.router.RegisterRequest(
		ctx,
		nodeID,
		requestingChainID,
		respondingChainID,
		requestID,
		op,
		failedMsg,
	)
}

func (r *tracedRouter) HandleInbound(ctx context.Context, msg message.InboundMessage) {
	m := msg.Message()
	destinationChainID, err := message.GetChainID(m)
	if err != nil {
		r.router.HandleInbound(ctx, msg)
		return
	}

	sourceChainID, err := message.GetSourceChainID(m)
	if err != nil {
		r.router.HandleInbound(ctx, msg)
		return
	}

	ctx, span := r.tracer.Start(ctx, "tracedRouter.HandleInbound", oteltrace.WithAttributes(
		attribute.Stringer("nodeID", msg.NodeID()),
		attribute.Stringer("messageOp", msg.Op()),
		attribute.Stringer("chainID", destinationChainID),
		attribute.Stringer("sourceChainID", sourceChainID),
	))
	defer span.End()

	r.router.HandleInbound(ctx, msg)
}

func (r *tracedRouter) Shutdown() {
	r.router.Shutdown()
}

func (r *tracedRouter) AddChain(chain handler.Handler) {
	r.router.AddChain(chain)
}

func (r *tracedRouter) Connected(nodeID ids.NodeID, nodeVersion *version.Application, subnetID ids.ID) {
	r.router.Connected(nodeID, nodeVersion, subnetID)
}

func (r *tracedRouter) Disconnected(nodeID ids.NodeID) {
	r.router.Disconnected(nodeID)
}

func (r *tracedRouter) Benched(chainID ids.ID, nodeID ids.NodeID) {
	r.router.Benched(chainID, nodeID)
}

func (r *tracedRouter) Unbenched(chainID ids.ID, nodeID ids.NodeID) {
	r.router.Unbenched(chainID, nodeID)
}

func (r *tracedRouter) HealthCheck() (interface{}, error) {
	return r.router.HealthCheck()
}
<|endoftext|>


<commit_before>package irckit

import (
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/mattermost/mattermost-server/model"
	"github.com/sorcix/irc"
)

// Channel is a representation of a room in our server
type Channel interface {
	Prefixer

	// ID is a normalized unique identifier for the channel
	ID() string

	// Created returns the time when the Channel was created.
	Created() time.Time

	// Names returns a sorted slice of Nicks in the channel
	Names() []string

	// Users returns a slice of Users in the channel.
	Users() []*User

	// HasUser returns whether a User is in the channel.
	HasUser(*User) bool

	// Invite prompts the User to join the Channel on behalf of Prefixer.
	Invite(from Prefixer, u *User) error

	// SendNamesResponse sends a User messages indicating the current members of the Channel.
	SendNamesResponse(u *User) error

	// Join introduces the User to the channel (handler for JOIN).
	Join(u *User) error

	// Part removes the User from the channel (handler for PART).
	Part(u *User, text string)

	// Message transmits a message from a User to the channel (handler for PRIVMSG).
	Message(u *User, text string)

	// Service returns the service that set the channel
	Service() string

	// Topic sets the topic of the channel (handler for TOPIC).
	Topic(from Prefixer, text string)

	// GetTopic gets the topic of the channel
	GetTopic() string

	// Unlink will disassociate the Channel from its Server.
	Unlink()

	// Len returns the number of Users in the channel.
	Len() int

	// String returns the name of the channel
	String() string

	// Spoof message
	SpoofMessage(from string, text string)

	// Spoof notice
	SpoofNotice(from string, text string)
}

type channel struct {
	created time.Time
	name    string
	server  Server
	id      string
	service string

	mu       sync.RWMutex
	topic    string
	usersIdx map[*User]struct{}
}

// NewChannel returns a Channel implementation for a given Server.
func NewChannel(server Server, channelId string, name string, service string) Channel {
	return &channel{
		created:  time.Now(),
		server:   server,
		id:       channelId,
		name:     name,
		service:  service,
		usersIdx: map[*User]struct{}{},
	}
}

func (ch *channel) GetTopic() string {
	return ch.topic
}

func (ch *channel) Prefix() *irc.Prefix {
	return ch.server.Prefix()
}

func (ch *channel) Service() string {
	return ch.service
}

func (ch *channel) String() string {
	return ch.name
}

// Created returns the time when the Channel was created.
func (ch *channel) Created() time.Time {
	return ch.created
}

// ID returns a normalized unique identifier for the channel.
func (ch *channel) ID() string {
	return ID(ch.id)
}

func (ch *channel) Message(from *User, text string) {
	for len(text) > 400 {
		msg := &irc.Message{
			Prefix:   from.Prefix(),
			Command:  irc.PRIVMSG,
			Params:   []string{ch.name},
			Trailing: text[:400] + "\n",
		}
		ch.mu.RLock()
		for to := range ch.usersIdx {
			to.Encode(msg)
		}
		ch.mu.RUnlock()
		text = text[400:]
	}
	msg := &irc.Message{
		Prefix:   from.Prefix(),
		Command:  irc.PRIVMSG,
		Params:   []string{ch.name},
		Trailing: text,
	}
	ch.mu.RLock()
	for to := range ch.usersIdx {
		to.Encode(msg)
	}
	ch.mu.RUnlock()
}

// Quit will remove the user from the channel and emit a PART message.
func (ch *channel) Part(u *User, text string) {
	msg := &irc.Message{
		Prefix:   u.Prefix(),
		Command:  irc.PART,
		Params:   []string{ch.name},
		Trailing: text,
	}
	ch.mu.Lock()
	if _, ok := ch.usersIdx[u]; !ok {
		ch.mu.Unlock()
		u.Encode(&irc.Message{
			Prefix:   ch.Prefix(),
			Command:  irc.ERR_NOTONCHANNEL,
			Params:   []string{ch.name},
			Trailing: "You're not on that channel",
		})
		return
	}
	u.Encode(msg)
	delete(ch.usersIdx, u)
	ch.mu.Unlock()
	u.Lock()
	delete(u.channels, ch)
	u.Unlock()

	for to := range ch.usersIdx {
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}
}

// Unlink will disassociate the Channel from the Server.
func (ch *channel) Unlink() {
	ch.server.UnlinkChannel(ch)
}

// Close will evict all users in the channel.
func (ch *channel) Close() error {
	ch.mu.Lock()
	for to := range ch.usersIdx {
		to.Encode(&irc.Message{
			Prefix:  to.Prefix(),
			Command: irc.PART,
			Params:  []string{ch.name},
		})
	}
	ch.usersIdx = map[*User]struct{}{}
	ch.mu.Unlock()
	return nil
}

// Invite prompts the User to join the Channel on behalf of Prefixer.
func (ch *channel) Invite(from Prefixer, u *User) error {
	return u.Encode(&irc.Message{
		Prefix:  from.Prefix(),
		Command: irc.INVITE,
		Params:  []string{u.Nick, ch.name},
	})
	// TODO: Save state that the user is invited?
}

// Topic sets the topic of the channel (handler for TOPIC).
func (ch *channel) Topic(from Prefixer, text string) {
	ch.mu.RLock()
	ch.topic = text
	// no newlines in topic
	ch.topic = strings.Replace(ch.topic, "\n", " ", -1)
	ch.topic = strings.Replace(ch.topic, "\r", " ", -1)

	msg := &irc.Message{
		Prefix:   from.Prefix(),
		Command:  irc.TOPIC,
		Params:   []string{ch.name},
		Trailing: ch.topic,
	}

	// only send join messages to real users
	for to := range ch.usersIdx {
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}

	ch.mu.RUnlock()
}

// SendNamesResponse sends a User messages indicating the current members of the Channel.
func (ch *channel) SendNamesResponse(u *User) error {
	msgs := []*irc.Message{}
	line := ""
	i := 0
	for _, name := range ch.Names() {
		if i+len(name) < 400 {
			line += name + " "
			i += len(name)
		} else {
			msgs = append(msgs, &irc.Message{
				Prefix:   ch.Prefix(),
				Command:  irc.RPL_NAMREPLY,
				Params:   []string{u.Nick, "=", ch.name},
				Trailing: line,
			})
			line = ""
			line += name + " "
			i = len(name)
		}
	}
	msgs = append(msgs, &irc.Message{
		Prefix:   ch.Prefix(),
		Command:  irc.RPL_NAMREPLY,
		Params:   []string{u.Nick, "=", ch.name},
		Trailing: line,
	})

	msgs = append(msgs, &irc.Message{
		Prefix:   ch.Prefix(),
		Params:   []string{u.Nick, ch.name},
		Command:  irc.RPL_ENDOFNAMES,
		Trailing: "End of /NAMES list.",
	})

	return u.Encode(msgs...)
}

// Join introduces a User to the channel (sends relevant messages, stores).
func (ch *channel) Join(u *User) error {
	// TODO: Check if user is already here?
	ch.mu.Lock()
	if _, exists := ch.usersIdx[u]; exists {
		ch.mu.Unlock()
		return nil
	}
	topic := ch.topic
	ch.usersIdx[u] = struct{}{}
	ch.mu.Unlock()
	u.Lock()
	u.channels[ch] = struct{}{}
	u.Unlock()

	// speed up &users join
	if ch.name == "&users" && u.MmGhostUser {
		return nil
	}

	msg := &irc.Message{
		Prefix:  u.Prefix(),
		Command: irc.JOIN,
		Params:  []string{ch.name},
	}

	// send regular users a notification of the join
	ch.mu.RLock()
	for to := range ch.usersIdx {
		// only send join messages to real users
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}
	ch.mu.RUnlock()

	msgs := []*irc.Message{}
	if topic != "" {
		msgs = append(msgs, &irc.Message{
			Prefix:   ch.Prefix(),
			Command:  irc.RPL_TOPIC,
			Params:   []string{u.Nick, ch.name},
			Trailing: topic,
		})
	}

	ch.SendNamesResponse(u)

	return u.Encode(msgs...)
}

func (ch *channel) HasUser(u *User) bool {
	ch.mu.RLock()
	_, ok := ch.usersIdx[u]
	ch.mu.RUnlock()
	return ok
}

// Users returns an unsorted slice of users who are in the channel.
func (ch *channel) Users() []*User {
	ch.mu.RLock()
	users := make([]*User, 0, len(ch.usersIdx))
	for u := range ch.usersIdx {
		users = append(users, u)
	}
	ch.mu.RUnlock()
	return users
}

// Names returns a sorted slice of Nick strings of users who are in the channel.
func (ch *channel) Names() []string {
	users := ch.Users()
	names := make([]string, 0, len(users))
	for _, u := range users {
		if strings.Contains(u.Roles, model.SYSTEM_ADMIN_ROLE_ID) {
			names = append(names, "@"+u.Nick)
		} else {
			names = append(names, u.Nick)
		}
	}
	// TODO: Append in sorted order?
	sort.Strings(names)
	return names
}

// Len returns the number of users in the channel.
func (ch *channel) Len() int {
	ch.mu.RLock()
	defer ch.mu.RUnlock()
	return len(ch.usersIdx)
}
<commit_msg>Unlock correctly. Fixes #246<commit_after>package irckit

import (
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/mattermost/mattermost-server/model"
	"github.com/sorcix/irc"
)

// Channel is a representation of a room in our server
type Channel interface {
	Prefixer

	// ID is a normalized unique identifier for the channel
	ID() string

	// Created returns the time when the Channel was created.
	Created() time.Time

	// Names returns a sorted slice of Nicks in the channel
	Names() []string

	// Users returns a slice of Users in the channel.
	Users() []*User

	// HasUser returns whether a User is in the channel.
	HasUser(*User) bool

	// Invite prompts the User to join the Channel on behalf of Prefixer.
	Invite(from Prefixer, u *User) error

	// SendNamesResponse sends a User messages indicating the current members of the Channel.
	SendNamesResponse(u *User) error

	// Join introduces the User to the channel (handler for JOIN).
	Join(u *User) error

	// Part removes the User from the channel (handler for PART).
	Part(u *User, text string)

	// Message transmits a message from a User to the channel (handler for PRIVMSG).
	Message(u *User, text string)

	// Service returns the service that set the channel
	Service() string

	// Topic sets the topic of the channel (handler for TOPIC).
	Topic(from Prefixer, text string)

	// GetTopic gets the topic of the channel
	GetTopic() string

	// Unlink will disassociate the Channel from its Server.
	Unlink()

	// Len returns the number of Users in the channel.
	Len() int

	// String returns the name of the channel
	String() string

	// Spoof message
	SpoofMessage(from string, text string)

	// Spoof notice
	SpoofNotice(from string, text string)
}

type channel struct {
	created time.Time
	name    string
	server  Server
	id      string
	service string

	mu       sync.RWMutex
	topic    string
	usersIdx map[*User]struct{}
}

// NewChannel returns a Channel implementation for a given Server.
func NewChannel(server Server, channelId string, name string, service string) Channel {
	return &channel{
		created:  time.Now(),
		server:   server,
		id:       channelId,
		name:     name,
		service:  service,
		usersIdx: map[*User]struct{}{},
	}
}

func (ch *channel) GetTopic() string {
	return ch.topic
}

func (ch *channel) Prefix() *irc.Prefix {
	return ch.server.Prefix()
}

func (ch *channel) Service() string {
	return ch.service
}

func (ch *channel) String() string {
	return ch.name
}

// Created returns the time when the Channel was created.
func (ch *channel) Created() time.Time {
	return ch.created
}

// ID returns a normalized unique identifier for the channel.
func (ch *channel) ID() string {
	return ID(ch.id)
}

func (ch *channel) Message(from *User, text string) {
	for len(text) > 400 {
		msg := &irc.Message{
			Prefix:   from.Prefix(),
			Command:  irc.PRIVMSG,
			Params:   []string{ch.name},
			Trailing: text[:400] + "\n",
		}
		ch.mu.RLock()
		for to := range ch.usersIdx {
			to.Encode(msg)
		}
		ch.mu.RUnlock()
		text = text[400:]
	}
	msg := &irc.Message{
		Prefix:   from.Prefix(),
		Command:  irc.PRIVMSG,
		Params:   []string{ch.name},
		Trailing: text,
	}
	ch.mu.RLock()
	for to := range ch.usersIdx {
		to.Encode(msg)
	}
	ch.mu.RUnlock()
}

// Quit will remove the user from the channel and emit a PART message.
func (ch *channel) Part(u *User, text string) {
	msg := &irc.Message{
		Prefix:   u.Prefix(),
		Command:  irc.PART,
		Params:   []string{ch.name},
		Trailing: text,
	}
	ch.mu.Lock()
	if _, ok := ch.usersIdx[u]; !ok {
		ch.mu.Unlock()
		u.Encode(&irc.Message{
			Prefix:   ch.Prefix(),
			Command:  irc.ERR_NOTONCHANNEL,
			Params:   []string{ch.name},
			Trailing: "You're not on that channel",
		})
		return
	}
	u.Encode(msg)
	delete(ch.usersIdx, u)
	u.Lock()
	delete(u.channels, ch)
	u.Unlock()

	for to := range ch.usersIdx {
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}
	ch.mu.Unlock()
}

// Unlink will disassociate the Channel from the Server.
func (ch *channel) Unlink() {
	ch.server.UnlinkChannel(ch)
}

// Close will evict all users in the channel.
func (ch *channel) Close() error {
	ch.mu.Lock()
	for to := range ch.usersIdx {
		to.Encode(&irc.Message{
			Prefix:  to.Prefix(),
			Command: irc.PART,
			Params:  []string{ch.name},
		})
	}
	ch.usersIdx = map[*User]struct{}{}
	ch.mu.Unlock()
	return nil
}

// Invite prompts the User to join the Channel on behalf of Prefixer.
func (ch *channel) Invite(from Prefixer, u *User) error {
	return u.Encode(&irc.Message{
		Prefix:  from.Prefix(),
		Command: irc.INVITE,
		Params:  []string{u.Nick, ch.name},
	})
	// TODO: Save state that the user is invited?
}

// Topic sets the topic of the channel (handler for TOPIC).
func (ch *channel) Topic(from Prefixer, text string) {
	ch.mu.RLock()
	ch.topic = text
	// no newlines in topic
	ch.topic = strings.Replace(ch.topic, "\n", " ", -1)
	ch.topic = strings.Replace(ch.topic, "\r", " ", -1)

	msg := &irc.Message{
		Prefix:   from.Prefix(),
		Command:  irc.TOPIC,
		Params:   []string{ch.name},
		Trailing: ch.topic,
	}

	// only send join messages to real users
	for to := range ch.usersIdx {
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}

	ch.mu.RUnlock()
}

// SendNamesResponse sends a User messages indicating the current members of the Channel.
func (ch *channel) SendNamesResponse(u *User) error {
	msgs := []*irc.Message{}
	line := ""
	i := 0
	for _, name := range ch.Names() {
		if i+len(name) < 400 {
			line += name + " "
			i += len(name)
		} else {
			msgs = append(msgs, &irc.Message{
				Prefix:   ch.Prefix(),
				Command:  irc.RPL_NAMREPLY,
				Params:   []string{u.Nick, "=", ch.name},
				Trailing: line,
			})
			line = ""
			line += name + " "
			i = len(name)
		}
	}
	msgs = append(msgs, &irc.Message{
		Prefix:   ch.Prefix(),
		Command:  irc.RPL_NAMREPLY,
		Params:   []string{u.Nick, "=", ch.name},
		Trailing: line,
	})

	msgs = append(msgs, &irc.Message{
		Prefix:   ch.Prefix(),
		Params:   []string{u.Nick, ch.name},
		Command:  irc.RPL_ENDOFNAMES,
		Trailing: "End of /NAMES list.",
	})

	return u.Encode(msgs...)
}

// Join introduces a User to the channel (sends relevant messages, stores).
func (ch *channel) Join(u *User) error {
	// TODO: Check if user is already here?
	ch.mu.Lock()
	if _, exists := ch.usersIdx[u]; exists {
		ch.mu.Unlock()
		return nil
	}
	topic := ch.topic
	ch.usersIdx[u] = struct{}{}
	ch.mu.Unlock()
	u.Lock()
	u.channels[ch] = struct{}{}
	u.Unlock()

	// speed up &users join
	if ch.name == "&users" && u.MmGhostUser {
		return nil
	}

	msg := &irc.Message{
		Prefix:  u.Prefix(),
		Command: irc.JOIN,
		Params:  []string{ch.name},
	}

	// send regular users a notification of the join
	ch.mu.RLock()
	for to := range ch.usersIdx {
		// only send join messages to real users
		if to.MmGhostUser == false {
			to.Encode(msg)
		}
	}
	ch.mu.RUnlock()

	msgs := []*irc.Message{}
	if topic != "" {
		msgs = append(msgs, &irc.Message{
			Prefix:   ch.Prefix(),
			Command:  irc.RPL_TOPIC,
			Params:   []string{u.Nick, ch.name},
			Trailing: topic,
		})
	}

	ch.SendNamesResponse(u)

	return u.Encode(msgs...)
}

func (ch *channel) HasUser(u *User) bool {
	ch.mu.RLock()
	_, ok := ch.usersIdx[u]
	ch.mu.RUnlock()
	return ok
}

// Users returns an unsorted slice of users who are in the channel.
func (ch *channel) Users() []*User {
	ch.mu.RLock()
	users := make([]*User, 0, len(ch.usersIdx))
	for u := range ch.usersIdx {
		users = append(users, u)
	}
	ch.mu.RUnlock()
	return users
}

// Names returns a sorted slice of Nick strings of users who are in the channel.
func (ch *channel) Names() []string {
	users := ch.Users()
	names := make([]string, 0, len(users))
	for _, u := range users {
		if strings.Contains(u.Roles, model.SYSTEM_ADMIN_ROLE_ID) {
			names = append(names, "@"+u.Nick)
		} else {
			names = append(names, u.Nick)
		}
	}
	// TODO: Append in sorted order?
	sort.Strings(names)
	return names
}

// Len returns the number of users in the channel.
func (ch *channel) Len() int {
	ch.mu.RLock()
	defer ch.mu.RUnlock()
	return len(ch.usersIdx)
}
<|endoftext|>


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

package mvcc

import (
	"log"
	"sync"
	"time"

	"github.com/coreos/etcd/lease"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/mvcc/mvccpb"
)

const (
	// chanBufLen is the length of the buffered chan
	// for sending out watched events.
	// TODO: find a good buf value. 1024 is just a random one that
	// seems to be reasonable.
	chanBufLen = 1024

	// maxWatchersPerSync is the number of watchers to sync in a single batch
	maxWatchersPerSync = 512
)

type watchable interface {
	watch(key, end []byte, startRev int64, id WatchID, ch chan<- WatchResponse) (*watcher, cancelFunc)
	progress(w *watcher)
	rev() int64
}

type watchableStore struct {
	mu sync.Mutex

	*store

	// contains all unsynced watchers that needs to sync with events that have happened
	unsynced watcherGroup

	// contains all synced watchers that are in sync with the progress of the store.
	// The key of the map is the key that the watcher watches on.
	synced watcherGroup

	stopc chan struct{}
	wg    sync.WaitGroup
}

// cancelFunc updates unsynced and synced maps when running
// cancel operations.
type cancelFunc func()

func New(b backend.Backend, le lease.Lessor, ig ConsistentIndexGetter) ConsistentWatchableKV {
	return newWatchableStore(b, le, ig)
}

func newWatchableStore(b backend.Backend, le lease.Lessor, ig ConsistentIndexGetter) *watchableStore {
	s := &watchableStore{
		store:    NewStore(b, le, ig),
		unsynced: newWatcherGroup(),
		synced:   newWatcherGroup(),
		stopc:    make(chan struct{}),
	}
	if s.le != nil {
		// use this store as the deleter so revokes trigger watch events
		s.le.SetRangeDeleter(s)
	}
	s.wg.Add(1)
	go s.syncWatchersLoop()
	return s
}

func (s *watchableStore) Put(key, value []byte, lease lease.LeaseID) (rev int64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	rev = s.store.Put(key, value, lease)
	changes := s.store.getChanges()
	if len(changes) != 1 {
		log.Panicf("unexpected len(changes) != 1 after put")
	}

	ev := mvccpb.Event{
		Type: mvccpb.PUT,
		Kv:   &changes[0],
	}
	s.notify(rev, []mvccpb.Event{ev})
	return rev
}

func (s *watchableStore) DeleteRange(key, end []byte) (n, rev int64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	n, rev = s.store.DeleteRange(key, end)
	changes := s.store.getChanges()

	if len(changes) != int(n) {
		log.Panicf("unexpected len(changes) != n after deleteRange")
	}

	if n == 0 {
		return n, rev
	}

	evs := make([]mvccpb.Event, n)
	for i := range changes {
		evs[i] = mvccpb.Event{
			Type: mvccpb.DELETE,
			Kv:   &changes[i]}
		evs[i].Kv.ModRevision = rev
	}
	s.notify(rev, evs)
	return n, rev
}

func (s *watchableStore) TxnBegin() int64 {
	s.mu.Lock()
	return s.store.TxnBegin()
}

func (s *watchableStore) TxnEnd(txnID int64) error {
	err := s.store.TxnEnd(txnID)
	if err != nil {
		return err
	}

	changes := s.getChanges()
	if len(changes) == 0 {
		s.mu.Unlock()
		return nil
	}

	rev := s.store.Rev()
	evs := make([]mvccpb.Event, len(changes))
	for i, change := range changes {
		switch change.CreateRevision {
		case 0:
			evs[i] = mvccpb.Event{
				Type: mvccpb.DELETE,
				Kv:   &changes[i]}
			evs[i].Kv.ModRevision = rev
		default:
			evs[i] = mvccpb.Event{
				Type: mvccpb.PUT,
				Kv:   &changes[i]}
		}
	}

	s.notify(rev, evs)
	s.mu.Unlock()

	return nil
}

func (s *watchableStore) Close() error {
	close(s.stopc)
	s.wg.Wait()
	return s.store.Close()
}

func (s *watchableStore) NewWatchStream() WatchStream {
	watchStreamGauge.Inc()
	return &watchStream{
		watchable: s,
		ch:        make(chan WatchResponse, chanBufLen),
		cancels:   make(map[WatchID]cancelFunc),
		watchers:  make(map[WatchID]*watcher),
	}
}

func (s *watchableStore) watch(key, end []byte, startRev int64, id WatchID, ch chan<- WatchResponse) (*watcher, cancelFunc) {
	s.mu.Lock()
	defer s.mu.Unlock()

	wa := &watcher{
		key: key,
		end: end,
		cur: startRev,
		id:  id,
		ch:  ch,
	}

	s.store.mu.Lock()
	synced := startRev > s.store.currentRev.main || startRev == 0
	if synced {
		wa.cur = s.store.currentRev.main + 1
		if startRev > wa.cur {
			wa.cur = startRev
		}
	}
	s.store.mu.Unlock()
	if synced {
		s.synced.add(wa)
	} else {
		slowWatcherGauge.Inc()
		s.unsynced.add(wa)
	}
	watcherGauge.Inc()

	cancel := cancelFunc(func() {
		s.mu.Lock()
		// remove references of the watcher
		if s.unsynced.delete(wa) {
			slowWatcherGauge.Dec()
			watcherGauge.Dec()
		} else if s.synced.delete(wa) {
			watcherGauge.Dec()
		}
		s.mu.Unlock()

		// If we cannot find it, it should have finished watch.
	})

	return wa, cancel
}

// syncWatchersLoop syncs the watcher in the unsynced map every 100ms.
func (s *watchableStore) syncWatchersLoop() {
	defer s.wg.Done()

	for {
		s.mu.Lock()
		st := time.Now()
		lastUnsyncedWatchers := s.unsynced.size()
		s.syncWatchers()
		unsyncedWatchers := s.unsynced.size()
		s.mu.Unlock()
		syncDuration := time.Since(st)

		waitDuration := 100 * time.Millisecond
		// more work pending?
		if unsyncedWatchers != 0 && lastUnsyncedWatchers > unsyncedWatchers {
			// be fair to other store operations by yielding time taken
			waitDuration = syncDuration
		}

		select {
		case <-time.After(waitDuration):
		case <-s.stopc:
			return
		}
	}
}

// syncWatchers syncs unsynced watchers by:
//	1. choose a set of watchers from the unsynced watcher group
//	2. iterate over the set to get the minimum revision and remove compacted watchers
//	3. use minimum revision to get all key-value pairs and send those events to watchers
//	4. remove synced watchers in set from unsynced group and move to synced group
func (s *watchableStore) syncWatchers() {
	if s.unsynced.size() == 0 {
		return
	}

	s.store.mu.Lock()
	defer s.store.mu.Unlock()

	// in order to find key-value pairs from unsynced watchers, we need to
	// find min revision index, and these revisions can be used to
	// query the backend store of key-value pairs
	curRev := s.store.currentRev.main
	compactionRev := s.store.compactMainRev
	wg, minRev := s.unsynced.choose(maxWatchersPerSync, curRev, compactionRev)
	minBytes, maxBytes := newRevBytes(), newRevBytes()
	revToBytes(revision{main: minRev}, minBytes)
	revToBytes(revision{main: curRev + 1}, maxBytes)

	// UnsafeRange returns keys and values. And in boltdb, keys are revisions.
	// values are actual key-value pairs in backend.
	tx := s.store.b.BatchTx()
	tx.Lock()
	revs, vs := tx.UnsafeRange(keyBucketName, minBytes, maxBytes, 0)
	evs := kvsToEvents(wg, revs, vs)
	tx.Unlock()

	wb := newWatcherBatch(wg, evs)
	for w := range wg.watchers {
		eb, ok := wb[w]
		if !ok {
			// bring un-notified watcher to synced
			w.cur = curRev
			s.synced.add(w)
			s.unsynced.delete(w)
			continue
		}

		select {
		case w.ch <- WatchResponse{WatchID: w.id, Events: eb.evs, Revision: curRev}:
			pendingEventsGauge.Add(float64(len(eb.evs)))
		default:
			// TODO: handle the full unsynced watchers.
			// continue to process other watchers for now, the full ones
			// will be processed next time and hopefully it will not be full.
			continue
		}
		if eb.moreRev != 0 {
			w.cur = eb.moreRev
			continue
		}
		w.cur = curRev
		s.synced.add(w)
		s.unsynced.delete(w)
	}

	slowWatcherGauge.Set(float64(s.unsynced.size()))
}

// kvsToEvents gets all events for the watchers from all key-value pairs
func kvsToEvents(wg *watcherGroup, revs, vals [][]byte) (evs []mvccpb.Event) {
	for i, v := range vals {
		var kv mvccpb.KeyValue
		if err := kv.Unmarshal(v); err != nil {
			log.Panicf("mvcc: cannot unmarshal event: %v", err)
		}

		if !wg.contains(string(kv.Key)) {
			continue
		}

		ty := mvccpb.PUT
		if isTombstone(revs[i]) {
			ty = mvccpb.DELETE
			// patch in mod revision so watchers won't skip
			kv.ModRevision = bytesToRev(revs[i]).main
		}
		evs = append(evs, mvccpb.Event{Kv: &kv, Type: ty})
	}
	return evs
}

// notify notifies the fact that given event at the given rev just happened to
// watchers that watch on the key of the event.
func (s *watchableStore) notify(rev int64, evs []mvccpb.Event) {
	for w, eb := range newWatcherBatch(&s.synced, evs) {
		if eb.revs != 1 {
			log.Panicf("mvcc: unexpected multiple revisions in notification")
		}
		select {
		case w.ch <- WatchResponse{WatchID: w.id, Events: eb.evs, Revision: s.Rev()}:
			pendingEventsGauge.Add(float64(len(eb.evs)))
		default:
			// move slow watcher to unsynced
			w.cur = rev
			s.unsynced.add(w)
			s.synced.delete(w)
			slowWatcherGauge.Inc()
		}
	}
}

func (s *watchableStore) rev() int64 { return s.store.Rev() }

func (s *watchableStore) progress(w *watcher) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if _, ok := s.synced.watchers[w]; ok {
		select {
		case w.ch <- WatchResponse{WatchID: w.id, Revision: s.rev()}:
		default:
			// If the ch is full, this watcher is receiving events.
			// We do not need to send progress at all.
		}
	}
}

type watcher struct {
	// the watcher key
	key []byte
	// end indicates the end of the range to watch.
	// If end is set, the watcher is on a range.
	end []byte

	// cur is the current watcher revision of a unsynced watcher.
	// cur will be updated for unsynced watcher while it is catching up.
	// cur is startRev of a synced watcher.
	// cur will not be updated for synced watcher.
	cur int64
	id  WatchID

	// a chan to send out the watch response.
	// The chan might be shared with other watchers.
	ch chan<- WatchResponse
}
<commit_msg>mvcc: move blocked sync watcher work to victim list<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mvcc

import (
	"log"
	"sync"
	"time"

	"github.com/coreos/etcd/lease"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/mvcc/mvccpb"
)

const (
	// chanBufLen is the length of the buffered chan
	// for sending out watched events.
	// TODO: find a good buf value. 1024 is just a random one that
	// seems to be reasonable.
	chanBufLen = 1024

	// maxWatchersPerSync is the number of watchers to sync in a single batch
	maxWatchersPerSync = 512
)

type watchable interface {
	watch(key, end []byte, startRev int64, id WatchID, ch chan<- WatchResponse) (*watcher, cancelFunc)
	progress(w *watcher)
	rev() int64
}

type watchableStore struct {
	mu sync.Mutex

	*store

	// victims are watcher batches that were blocked on the watch channel
	victims []watcherBatch
	victimc chan struct{}

	// contains all unsynced watchers that needs to sync with events that have happened
	unsynced watcherGroup

	// contains all synced watchers that are in sync with the progress of the store.
	// The key of the map is the key that the watcher watches on.
	synced watcherGroup

	stopc chan struct{}
	wg    sync.WaitGroup
}

// cancelFunc updates unsynced and synced maps when running
// cancel operations.
type cancelFunc func()

func New(b backend.Backend, le lease.Lessor, ig ConsistentIndexGetter) ConsistentWatchableKV {
	return newWatchableStore(b, le, ig)
}

func newWatchableStore(b backend.Backend, le lease.Lessor, ig ConsistentIndexGetter) *watchableStore {
	s := &watchableStore{
		store:    NewStore(b, le, ig),
		victimc:  make(chan struct{}, 1),
		unsynced: newWatcherGroup(),
		synced:   newWatcherGroup(),
		stopc:    make(chan struct{}),
	}
	if s.le != nil {
		// use this store as the deleter so revokes trigger watch events
		s.le.SetRangeDeleter(s)
	}
	s.wg.Add(2)
	go s.syncWatchersLoop()
	go s.syncVictimsLoop()
	return s
}

func (s *watchableStore) Put(key, value []byte, lease lease.LeaseID) (rev int64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	rev = s.store.Put(key, value, lease)
	changes := s.store.getChanges()
	if len(changes) != 1 {
		log.Panicf("unexpected len(changes) != 1 after put")
	}

	ev := mvccpb.Event{
		Type: mvccpb.PUT,
		Kv:   &changes[0],
	}
	s.notify(rev, []mvccpb.Event{ev})
	return rev
}

func (s *watchableStore) DeleteRange(key, end []byte) (n, rev int64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	n, rev = s.store.DeleteRange(key, end)
	changes := s.store.getChanges()

	if len(changes) != int(n) {
		log.Panicf("unexpected len(changes) != n after deleteRange")
	}

	if n == 0 {
		return n, rev
	}

	evs := make([]mvccpb.Event, n)
	for i := range changes {
		evs[i] = mvccpb.Event{
			Type: mvccpb.DELETE,
			Kv:   &changes[i]}
		evs[i].Kv.ModRevision = rev
	}
	s.notify(rev, evs)
	return n, rev
}

func (s *watchableStore) TxnBegin() int64 {
	s.mu.Lock()
	return s.store.TxnBegin()
}

func (s *watchableStore) TxnEnd(txnID int64) error {
	err := s.store.TxnEnd(txnID)
	if err != nil {
		return err
	}

	changes := s.getChanges()
	if len(changes) == 0 {
		s.mu.Unlock()
		return nil
	}

	rev := s.store.Rev()
	evs := make([]mvccpb.Event, len(changes))
	for i, change := range changes {
		switch change.CreateRevision {
		case 0:
			evs[i] = mvccpb.Event{
				Type: mvccpb.DELETE,
				Kv:   &changes[i]}
			evs[i].Kv.ModRevision = rev
		default:
			evs[i] = mvccpb.Event{
				Type: mvccpb.PUT,
				Kv:   &changes[i]}
		}
	}

	s.notify(rev, evs)
	s.mu.Unlock()

	return nil
}

func (s *watchableStore) Close() error {
	close(s.stopc)
	s.wg.Wait()
	return s.store.Close()
}

func (s *watchableStore) NewWatchStream() WatchStream {
	watchStreamGauge.Inc()
	return &watchStream{
		watchable: s,
		ch:        make(chan WatchResponse, chanBufLen),
		cancels:   make(map[WatchID]cancelFunc),
		watchers:  make(map[WatchID]*watcher),
	}
}

func (s *watchableStore) watch(key, end []byte, startRev int64, id WatchID, ch chan<- WatchResponse) (*watcher, cancelFunc) {
	s.mu.Lock()
	defer s.mu.Unlock()

	wa := &watcher{
		key: key,
		end: end,
		cur: startRev,
		id:  id,
		ch:  ch,
	}

	s.store.mu.Lock()
	synced := startRev > s.store.currentRev.main || startRev == 0
	if synced {
		wa.cur = s.store.currentRev.main + 1
		if startRev > wa.cur {
			wa.cur = startRev
		}
	}
	s.store.mu.Unlock()
	if synced {
		s.synced.add(wa)
	} else {
		slowWatcherGauge.Inc()
		s.unsynced.add(wa)
	}
	watcherGauge.Inc()

	cancel := cancelFunc(func() {
		s.mu.Lock()
		// remove references of the watcher
		if s.unsynced.delete(wa) {
			slowWatcherGauge.Dec()
			watcherGauge.Dec()
		} else if s.synced.delete(wa) {
			watcherGauge.Dec()
		} else {
			for _, wb := range s.victims {
				if wb[wa] != nil {
					slowWatcherGauge.Dec()
					watcherGauge.Dec()
					delete(wb, wa)
					break
				}
			}
		}
		s.mu.Unlock()

		// If we cannot find it, it should have finished watch.
	})

	return wa, cancel
}

// syncWatchersLoop syncs the watcher in the unsynced map every 100ms.
func (s *watchableStore) syncWatchersLoop() {
	defer s.wg.Done()

	for {
		s.mu.Lock()
		st := time.Now()
		lastUnsyncedWatchers := s.unsynced.size()
		s.syncWatchers()
		unsyncedWatchers := s.unsynced.size()
		s.mu.Unlock()
		syncDuration := time.Since(st)

		waitDuration := 100 * time.Millisecond
		// more work pending?
		if unsyncedWatchers != 0 && lastUnsyncedWatchers > unsyncedWatchers {
			// be fair to other store operations by yielding time taken
			waitDuration = syncDuration
		}

		select {
		case <-time.After(waitDuration):
		case <-s.stopc:
			return
		}
	}
}

// syncVictimsLoop tries to write precomputed watcher responses to
// watchers that had a blocked watcher channel
func (s *watchableStore) syncVictimsLoop() {
	defer s.wg.Done()

	for {
		for s.moveVictims() != 0 {
			// try to update all victim watchers
		}
		s.mu.Lock()
		isEmpty := len(s.victims) == 0
		s.mu.Unlock()

		var tickc <-chan time.Time
		if !isEmpty {
			tickc = time.After(10 * time.Millisecond)
		}

		select {
		case <-tickc:
		case <-s.victimc:
		case <-s.stopc:
			return
		}
	}
}

// moveVictims tries to update watches with already pending event data
func (s *watchableStore) moveVictims() (moved int) {
	s.mu.Lock()
	victims := s.victims
	s.victims = nil
	s.mu.Unlock()

	var newVictim watcherBatch
	for _, wb := range victims {
		// try to send responses again
		for w, eb := range wb {
			select {
			case w.ch <- WatchResponse{WatchID: w.id, Events: eb.evs, Revision: w.cur}:
				pendingEventsGauge.Add(float64(len(eb.evs)))
			default:
				if newVictim == nil {
					newVictim = make(watcherBatch)
				}
				newVictim[w] = eb
				continue
			}
			moved++
		}

		// assign completed victim watchers to unsync/sync
		s.mu.Lock()
		s.store.mu.Lock()
		curRev := s.store.currentRev.main
		for w, eb := range wb {
			if newVictim != nil && newVictim[w] != nil {
				// couldn't send watch response; stays victim
				continue
			}
			if eb.moreRev != 0 {
				w.cur = eb.moreRev
			}
			if w.cur < curRev {
				s.unsynced.add(w)
			} else {
				slowWatcherGauge.Dec()
				s.synced.add(w)
			}
		}
		s.store.mu.Unlock()
		s.mu.Unlock()
	}

	if len(newVictim) > 0 {
		s.mu.Lock()
		s.victims = append(s.victims, newVictim)
		s.mu.Unlock()
	}

	return moved
}

// syncWatchers syncs unsynced watchers by:
//	1. choose a set of watchers from the unsynced watcher group
//	2. iterate over the set to get the minimum revision and remove compacted watchers
//	3. use minimum revision to get all key-value pairs and send those events to watchers
//	4. remove synced watchers in set from unsynced group and move to synced group
func (s *watchableStore) syncWatchers() {
	if s.unsynced.size() == 0 {
		return
	}

	s.store.mu.Lock()
	defer s.store.mu.Unlock()

	// in order to find key-value pairs from unsynced watchers, we need to
	// find min revision index, and these revisions can be used to
	// query the backend store of key-value pairs
	curRev := s.store.currentRev.main
	compactionRev := s.store.compactMainRev
	wg, minRev := s.unsynced.choose(maxWatchersPerSync, curRev, compactionRev)
	minBytes, maxBytes := newRevBytes(), newRevBytes()
	revToBytes(revision{main: minRev}, minBytes)
	revToBytes(revision{main: curRev + 1}, maxBytes)

	// UnsafeRange returns keys and values. And in boltdb, keys are revisions.
	// values are actual key-value pairs in backend.
	tx := s.store.b.BatchTx()
	tx.Lock()
	revs, vs := tx.UnsafeRange(keyBucketName, minBytes, maxBytes, 0)
	evs := kvsToEvents(wg, revs, vs)
	tx.Unlock()

	wb := newWatcherBatch(wg, evs)
	for w := range wg.watchers {
		eb, ok := wb[w]
		if !ok {
			// bring un-notified watcher to synced
			w.cur = curRev
			s.synced.add(w)
			s.unsynced.delete(w)
			continue
		}

		select {
		case w.ch <- WatchResponse{WatchID: w.id, Events: eb.evs, Revision: curRev}:
			pendingEventsGauge.Add(float64(len(eb.evs)))
		default:
			// TODO: handle the full unsynced watchers.
			// continue to process other watchers for now, the full ones
			// will be processed next time and hopefully it will not be full.
			continue
		}
		if eb.moreRev != 0 {
			w.cur = eb.moreRev
			continue
		}
		w.cur = curRev
		s.synced.add(w)
		s.unsynced.delete(w)
	}

	vsz := 0
	for _, v := range s.victims {
		vsz += len(v)
	}
	slowWatcherGauge.Set(float64(s.unsynced.size() + vsz))
}

// kvsToEvents gets all events for the watchers from all key-value pairs
func kvsToEvents(wg *watcherGroup, revs, vals [][]byte) (evs []mvccpb.Event) {
	for i, v := range vals {
		var kv mvccpb.KeyValue
		if err := kv.Unmarshal(v); err != nil {
			log.Panicf("mvcc: cannot unmarshal event: %v", err)
		}

		if !wg.contains(string(kv.Key)) {
			continue
		}

		ty := mvccpb.PUT
		if isTombstone(revs[i]) {
			ty = mvccpb.DELETE
			// patch in mod revision so watchers won't skip
			kv.ModRevision = bytesToRev(revs[i]).main
		}
		evs = append(evs, mvccpb.Event{Kv: &kv, Type: ty})
	}
	return evs
}

// notify notifies the fact that given event at the given rev just happened to
// watchers that watch on the key of the event.
func (s *watchableStore) notify(rev int64, evs []mvccpb.Event) {
	var victim watcherBatch
	for w, eb := range newWatcherBatch(&s.synced, evs) {
		if eb.revs != 1 {
			log.Panicf("mvcc: unexpected multiple revisions in notification")
		}
		select {
		case w.ch <- WatchResponse{WatchID: w.id, Events: eb.evs, Revision: s.Rev()}:
			pendingEventsGauge.Add(float64(len(eb.evs)))
		default:
			// move slow watcher to victims
			w.cur = rev
			if victim == nil {
				victim = make(watcherBatch)
			}
			victim[w] = eb
			s.synced.delete(w)
			slowWatcherGauge.Inc()
		}
	}
	s.addVictim(victim)
}

func (s *watchableStore) addVictim(victim watcherBatch) {
	if victim == nil {
		return
	}
	s.victims = append(s.victims, victim)
	select {
	case s.victimc <- struct{}{}:
	default:
	}
}

func (s *watchableStore) rev() int64 { return s.store.Rev() }

func (s *watchableStore) progress(w *watcher) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if _, ok := s.synced.watchers[w]; ok {
		select {
		case w.ch <- WatchResponse{WatchID: w.id, Revision: s.rev()}:
		default:
			// If the ch is full, this watcher is receiving events.
			// We do not need to send progress at all.
		}
	}
}

type watcher struct {
	// the watcher key
	key []byte
	// end indicates the end of the range to watch.
	// If end is set, the watcher is on a range.
	end []byte

	// cur is the current watcher revision of a unsynced watcher.
	// cur will be updated for unsynced watcher while it is catching up.
	// cur is startRev of a synced watcher.
	// cur will not be updated for synced watcher.
	cur int64
	id  WatchID

	// a chan to send out the watch response.
	// The chan might be shared with other watchers.
	ch chan<- WatchResponse
}
<|endoftext|>


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

package namespace

import (
	"context"
	"fmt"
	"strings"
	"sync"

	eh "github.com/looplab/eventhorizon"
	"github.com/looplab/eventhorizon/uuid"
)

// EventStore is an event store with support for namespaces passed in the context.
type EventStore struct {
	eventStores   map[string]eh.EventStore
	eventStoresMu sync.RWMutex
	newEventStore func(ns string) (eh.EventStore, error)
}

// NewEventStore creates a new event store which will use the provided factory
// function to create new event stores for the provided namespace.
//
// Usage:
//    eventStore := NewEventStore(func(ns string) (eh.EventStore, error) {
//        s, err := mongodb.NewEventStore("mongodb://", ns)
//        if err != nil {
//            return nil, err
//        }
//        return s, nil
//    })
//
// Usage shared DB client:
//    client, err := mongo.Connect(ctx)
//    ...
//
//    eventStore := NewEventStore(func(ns string) (eh.EventStore, error) {
//        s, err := mongodb.NewEventStoreWithClient(client, ns)
//        if err != nil {
//            return nil, err
//        }
//        return s, nil
//    })
func NewEventStore(factory func(ns string) (eh.EventStore, error)) *EventStore {
	return &EventStore{
		eventStores:   map[string]eh.EventStore{},
		newEventStore: factory,
	}
}

// PreRegisterNamespace will make sure that a namespace exists in the eventstore.
// In normal cases the eventstore for a namespace is created when an event for
// that namespace is first seen.
func (s *EventStore) PreRegisterNamespace(ns string) error {
	ctx := NewContext(context.Background(), ns)

	// This creates the eventstore for a namespace in case id did not exist.
	_, err := s.eventStore(ctx)

	return err
}

// Save implements the Save method of the eventhorizon.EventStore interface.
func (s *EventStore) Save(ctx context.Context, events []eh.Event, originalVersion int) error {
	store, err := s.eventStore(ctx)
	if err != nil {
		return err
	}

	return store.Save(ctx, events, originalVersion)
}

// Load implements the Load method of the eventhorizon.EventStore interface.
func (s *EventStore) Load(ctx context.Context, id uuid.UUID) ([]eh.Event, error) {
	store, err := s.eventStore(ctx)
	if err != nil {
		return nil, err
	}

	return store.Load(ctx, id)
}

// Close implements the Close method of the eventhorizon.EventStore interface.
func (s *EventStore) Close() error {
	s.eventStoresMu.RLock()
	defer s.eventStoresMu.RUnlock()

	var errStrs []string

	for _, store := range s.eventStores {
		if err := store.Close(); err != nil {
			errStrs = append(errStrs, err.Error())
		}
	}

	if len(errStrs) > 0 {
		return fmt.Errorf("multiple errors: %s", strings.Join(errStrs, ", "))
	}

	return nil
}

// eventStore is a helper that returns or creates event stores for each namespace.
func (s *EventStore) eventStore(ctx context.Context) (eh.EventStore, error) {
	ns := FromContext(ctx)

	s.eventStoresMu.RLock()
	eventStore, ok := s.eventStores[ns]
	s.eventStoresMu.RUnlock()

	if !ok {
		s.eventStoresMu.Lock()
		defer s.eventStoresMu.Unlock()

		// Perform an additional existence check within the write lock in the
		// unlikely event that someone else added the event store right before us.
		if _, ok := s.eventStores[ns]; !ok {
			var err error

			eventStore, err = s.newEventStore(ns)
			if err != nil {
				return nil, fmt.Errorf("could not create event store for namespace '%s': %w", ns, err)
			}

			s.eventStores[ns] = eventStore
		}
	}

	return eventStore, nil
}
<commit_msg>set eventStore variable in the inner check<commit_after>// Copyright (c) 2021 - The Event Horizon authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package namespace

import (
	"context"
	"fmt"
	"strings"
	"sync"

	eh "github.com/looplab/eventhorizon"
	"github.com/looplab/eventhorizon/uuid"
)

// EventStore is an event store with support for namespaces passed in the context.
type EventStore struct {
	eventStores   map[string]eh.EventStore
	eventStoresMu sync.RWMutex
	newEventStore func(ns string) (eh.EventStore, error)
}

// NewEventStore creates a new event store which will use the provided factory
// function to create new event stores for the provided namespace.
//
// Usage:
//    eventStore := NewEventStore(func(ns string) (eh.EventStore, error) {
//        s, err := mongodb.NewEventStore("mongodb://", ns)
//        if err != nil {
//            return nil, err
//        }
//        return s, nil
//    })
//
// Usage shared DB client:
//    client, err := mongo.Connect(ctx)
//    ...
//
//    eventStore := NewEventStore(func(ns string) (eh.EventStore, error) {
//        s, err := mongodb.NewEventStoreWithClient(client, ns)
//        if err != nil {
//            return nil, err
//        }
//        return s, nil
//    })
func NewEventStore(factory func(ns string) (eh.EventStore, error)) *EventStore {
	return &EventStore{
		eventStores:   map[string]eh.EventStore{},
		newEventStore: factory,
	}
}

// PreRegisterNamespace will make sure that a namespace exists in the eventstore.
// In normal cases the eventstore for a namespace is created when an event for
// that namespace is first seen.
func (s *EventStore) PreRegisterNamespace(ns string) error {
	ctx := NewContext(context.Background(), ns)

	// This creates the eventstore for a namespace in case id did not exist.
	_, err := s.eventStore(ctx)

	return err
}

// Save implements the Save method of the eventhorizon.EventStore interface.
func (s *EventStore) Save(ctx context.Context, events []eh.Event, originalVersion int) error {
	store, err := s.eventStore(ctx)
	if err != nil {
		return err
	}

	return store.Save(ctx, events, originalVersion)
}

// Load implements the Load method of the eventhorizon.EventStore interface.
func (s *EventStore) Load(ctx context.Context, id uuid.UUID) ([]eh.Event, error) {
	store, err := s.eventStore(ctx)
	if err != nil {
		return nil, err
	}

	return store.Load(ctx, id)
}

// Close implements the Close method of the eventhorizon.EventStore interface.
func (s *EventStore) Close() error {
	s.eventStoresMu.RLock()
	defer s.eventStoresMu.RUnlock()

	var errStrs []string

	for _, store := range s.eventStores {
		if err := store.Close(); err != nil {
			errStrs = append(errStrs, err.Error())
		}
	}

	if len(errStrs) > 0 {
		return fmt.Errorf("multiple errors: %s", strings.Join(errStrs, ", "))
	}

	return nil
}

// eventStore is a helper that returns or creates event stores for each namespace.
func (s *EventStore) eventStore(ctx context.Context) (eh.EventStore, error) {
	ns := FromContext(ctx)

	s.eventStoresMu.RLock()
	eventStore, ok := s.eventStores[ns]
	s.eventStoresMu.RUnlock()

	if !ok {
		s.eventStoresMu.Lock()
		defer s.eventStoresMu.Unlock()

		// Perform an additional existence check within the write lock in the
		// unlikely event that someone else added the event store right before us.
		if eventStore, ok = s.eventStores[ns]; !ok {
			var err error

			eventStore, err = s.newEventStore(ns)
			if err != nil {
				return nil, fmt.Errorf("could not create event store for namespace '%s': %w", ns, err)
			}

			s.eventStores[ns] = eventStore
		}
	}

	return eventStore, nil
}
<|endoftext|>