<commit_before>// +build integration

package ngalert

import (
	"encoding/json"
	"errors"
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/registry"
	"github.com/grafana/grafana/pkg/services/ngalert/eval"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func mockTimeNow() {
	var timeSeed int64
	timeNow = func() time.Time {
		fakeNow := time.Unix(timeSeed, 0).UTC()
		timeSeed++
		return fakeNow
	}
}

func resetTimeNow() {
	timeNow = time.Now
}

func TestCreatingAlertDefinition(t *testing.T) {
	mockTimeNow()
	defer resetTimeNow()

	var customIntervalSeconds int64 = 120
	testCases := []struct {
		desc                 string
		inputIntervalSeconds *int64
		inputTitle           string
		expectedError        error
		expectedInterval     int64

		expectedUpdated time.Time
	}{
		{
			desc:                 "should create successfuly an alert definition with default interval",
			inputIntervalSeconds: nil,
			inputTitle:           "a name",
			expectedInterval:     defaultIntervalSeconds,
			expectedUpdated:      time.Unix(0, 0).UTC(),
		},
		{
			desc:                 "should create successfuly an alert definition with custom interval",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           "another name",
			expectedInterval:     customIntervalSeconds,
			expectedUpdated:      time.Unix(1, 0).UTC(),
		},
		{
			desc:                 "should fail to create an alert definition with too big name",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           getLongString(alertDefinitionMaxTitleLength + 1),
			expectedError:        errors.New(""),
		},
		{
			desc:                 "should fail to create an alert definition with empty title",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           "",
			expectedError:        errEmptyTitleError,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			ng := setupTestEnv(t)
			t.Cleanup(registry.ClearOverrides)

			q := saveAlertDefinitionCommand{
				OrgID:     1,
				Title:     tc.inputTitle,
				Condition: "B",
				Data: []eval.AlertQuery{
					{
						Model: json.RawMessage(`{
								"datasource": "__expr__",
								"type":"math",
								"expression":"2 + 3 > 1"
							}`),
						RefID: "B",
						RelativeTimeRange: eval.RelativeTimeRange{
							From: eval.Duration(time.Duration(5) * time.Hour),
							To:   eval.Duration(time.Duration(3) * time.Hour),
						},
					},
				},
			}
			if tc.inputIntervalSeconds != nil {
				q.IntervalSeconds = tc.inputIntervalSeconds
			}
			err := ng.saveAlertDefinition(&q)
			switch {
			case tc.expectedError != nil:
				require.Error(t, err)
			default:
				require.NoError(t, err)
				assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
				assert.Equal(t, tc.expectedInterval, q.Result.IntervalSeconds)
				assert.Equal(t, int64(1), q.Result.Version)

			}
		})
	}
}
func TestCreatingConflictionAlertDefinition(t *testing.T) {
	t.Run("Should fail to create alert definition with conflicting org_id, title", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := saveAlertDefinitionCommand{
			OrgID:     1,
			Title:     "title",
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
								"datasource": "__expr__",
								"type":"math",
								"expression":"2 + 3 > 1"
							}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(time.Duration(5) * time.Hour),
						To:   eval.Duration(time.Duration(3) * time.Hour),
					},
				},
			},
		}

		err := ng.saveAlertDefinition(&q)
		require.NoError(t, err)

		err = ng.saveAlertDefinition(&q)
		require.Error(t, err)
		assert.True(t, ng.SQLStore.Dialect.IsUniqueConstraintViolation(err))
	})
}

func TestUpdatingAlertDefinition(t *testing.T) {
	t.Run("zero rows affected when updating unknown alert", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := updateAlertDefinitionCommand{
			UID:       "unknown",
			OrgID:     1,
			Title:     "something completely different",
			Condition: "A",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 2 > 1"
						}`),
					RefID: "A",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(time.Duration(5) * time.Hour),
						To:   eval.Duration(time.Duration(3) * time.Hour),
					},
				},
			},
		}

		err := ng.updateAlertDefinition(&q)
		require.NoError(t, err)
	})

	t.Run("updating existing alert", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		var initialInterval int64 = 120
		alertDefinition := createTestAlertDefinition(t, ng, initialInterval)
		created := alertDefinition.Updated

		var customInterval int64 = 30
		testCases := []struct {
			desc                    string
			inputOrgID              int64
			inputTitle              string
			inputInterval           *int64
			expectedError           error
			expectedIntervalSeconds int64
			expectedUpdated         time.Time
			expectedTitle           string
		}{
			{
				desc:                    "should not update previous interval if it's not provided",
				inputInterval:           nil,
				inputOrgID:              alertDefinition.OrgID,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: initialInterval,
				expectedUpdated:         time.Unix(1, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:                    "should update interval if it's provided",
				inputInterval:           &customInterval,
				inputOrgID:              alertDefinition.OrgID,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(2, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:                    "should not update organisation if it's provided",
				inputInterval:           &customInterval,
				inputOrgID:              0,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(3, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:          "should not update alert definition if the title it's too big",
				inputInterval: &customInterval,
				inputOrgID:    0,
				inputTitle:    getLongString(alertDefinitionMaxTitleLength + 1),
				expectedError: errors.New(""),
			},
			{
				desc:                    "should not update alert definition title if the title is empty",
				inputInterval:           &customInterval,
				inputOrgID:              0,
				inputTitle:              "",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(4, 0).UTC(),
				expectedTitle:           "something completely different",
			},
		}

		q := updateAlertDefinitionCommand{
			UID:       (*alertDefinition).UID,
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 3 > 1"
						}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(5 * time.Hour),
						To:   eval.Duration(3 * time.Hour),
					},
				},
			},
		}

		lastUpdated := created
		previousAlertDefinition := alertDefinition
		for _, tc := range testCases {
			t.Run(tc.desc, func(t *testing.T) {
				if tc.inputInterval != nil {
					q.IntervalSeconds = tc.inputInterval
				}
				if tc.inputOrgID != 0 {
					q.OrgID = tc.inputOrgID
				}
				q.Title = tc.inputTitle
				err := ng.updateAlertDefinition(&q)
				switch {
				case tc.expectedError != nil:
					require.Error(t, err)

					assert.Equal(t, previousAlertDefinition.Title, q.Result.Title)
					assert.Equal(t, previousAlertDefinition.Condition, q.Result.Condition)
					assert.Equal(t, len(previousAlertDefinition.Data), len(q.Result.Data))
					assert.Equal(t, previousAlertDefinition.IntervalSeconds, q.Result.IntervalSeconds)
					assert.Equal(t, previousAlertDefinition.Updated, q.Result.Updated)
					assert.Equal(t, previousAlertDefinition.Version, q.Result.Version)
					assert.Equal(t, previousAlertDefinition.OrgID, q.Result.OrgID)
					assert.Equal(t, previousAlertDefinition.UID, q.Result.UID)
				default:
					require.NoError(t, err)
					assert.Equal(t, previousAlertDefinition.ID, q.Result.ID)
					assert.Equal(t, previousAlertDefinition.UID, q.Result.UID)
					assert.True(t, q.Result.Updated.After(lastUpdated))
					assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
					assert.Equal(t, previousAlertDefinition.Version+1, q.Result.Version)

					assert.Equal(t, alertDefinition.OrgID, q.Result.OrgID)

					assert.Equal(t, "something completely different", q.Result.Title)
					assert.Equal(t, "B", q.Result.Condition)
					assert.Equal(t, 1, len(q.Result.Data))
					assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
					assert.Equal(t, tc.expectedIntervalSeconds, q.Result.IntervalSeconds)
					assert.Equal(t, previousAlertDefinition.Version+1, q.Result.Version)
					assert.Equal(t, alertDefinition.OrgID, q.Result.OrgID)
					assert.Equal(t, alertDefinition.UID, q.Result.UID)

					previousAlertDefinition = q.Result
				}
			})

		}

	})
}

func TestUpdatingConflictingAlertDefinition(t *testing.T) {
	t.Run("should fail to update alert definition with reserved title", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		var initialInterval int64 = 120
		alertDef1 := createTestAlertDefinition(t, ng, initialInterval)
		alertDef2 := createTestAlertDefinition(t, ng, initialInterval)

		q := updateAlertDefinitionCommand{
			UID:       (*alertDef2).UID,
			Title:     alertDef1.Title,
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 3 > 1"
						}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(5 * time.Hour),
						To:   eval.Duration(3 * time.Hour),
					},
				},
			},
		}

		err := ng.updateAlertDefinition(&q)
		require.Error(t, err)
		assert.True(t, ng.SQLStore.Dialect.IsUniqueConstraintViolation(err))
	})
}

func TestDeletingAlertDefinition(t *testing.T) {
	t.Run("zero rows affected when deleting unknown alert", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := deleteAlertDefinitionByUIDCommand{
			UID:   "unknown",
			OrgID: 1,
		}

		err := ng.deleteAlertDefinitionByUID(&q)
		require.NoError(t, err)
	})

	t.Run("deleting successfully existing alert", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		alertDefinition := createTestAlertDefinition(t, ng, 60)

		q := deleteAlertDefinitionByUIDCommand{
			UID:   (*alertDefinition).UID,
			OrgID: 1,
		}

		// save an instance for the definition
		saveCmd := &saveAlertInstanceCommand{
			DefinitionOrgID: alertDefinition.OrgID,
			DefinitionUID:   alertDefinition.UID,
			State:           InstanceStateFiring,
			Labels:          InstanceLabels{"test": "testValue"},
		}
		err := ng.saveAlertInstance(saveCmd)
		require.NoError(t, err)
		listCommand := &listAlertInstancesQuery{
			DefinitionOrgID: alertDefinition.OrgID,
			DefinitionUID:   alertDefinition.UID,
		}
		err = ng.listAlertInstances(listCommand)
		require.NoError(t, err)
		require.Len(t, listCommand.Result, 1)

		err = ng.deleteAlertDefinitionByUID(&q)
		require.NoError(t, err)

		// assert that alert instance is deleted
		err = ng.listAlertInstances(listCommand)
		require.NoError(t, err)

		require.Len(t, listCommand.Result, 0)
	})
}

func getLongString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = 'a'
	}
	return string(b)
}
<commit_msg>Chore: Fixes small typos (#31461)<commit_after>// +build integration

package ngalert

import (
	"encoding/json"
	"errors"
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/registry"
	"github.com/grafana/grafana/pkg/services/ngalert/eval"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func mockTimeNow() {
	var timeSeed int64
	timeNow = func() time.Time {
		fakeNow := time.Unix(timeSeed, 0).UTC()
		timeSeed++
		return fakeNow
	}
}

func resetTimeNow() {
	timeNow = time.Now
}

func TestCreatingAlertDefinition(t *testing.T) {
	mockTimeNow()
	defer resetTimeNow()

	var customIntervalSeconds int64 = 120
	testCases := []struct {
		desc                 string
		inputIntervalSeconds *int64
		inputTitle           string
		expectedError        error
		expectedInterval     int64

		expectedUpdated time.Time
	}{
		{
			desc:                 "should create successfully an alert definition with default interval",
			inputIntervalSeconds: nil,
			inputTitle:           "a name",
			expectedInterval:     defaultIntervalSeconds,
			expectedUpdated:      time.Unix(0, 0).UTC(),
		},
		{
			desc:                 "should create successfully an alert definition with custom interval",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           "another name",
			expectedInterval:     customIntervalSeconds,
			expectedUpdated:      time.Unix(1, 0).UTC(),
		},
		{
			desc:                 "should fail to create an alert definition with too big name",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           getLongString(alertDefinitionMaxTitleLength + 1),
			expectedError:        errors.New(""),
		},
		{
			desc:                 "should fail to create an alert definition with empty title",
			inputIntervalSeconds: &customIntervalSeconds,
			inputTitle:           "",
			expectedError:        errEmptyTitleError,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			ng := setupTestEnv(t)
			t.Cleanup(registry.ClearOverrides)

			q := saveAlertDefinitionCommand{
				OrgID:     1,
				Title:     tc.inputTitle,
				Condition: "B",
				Data: []eval.AlertQuery{
					{
						Model: json.RawMessage(`{
								"datasource": "__expr__",
								"type":"math",
								"expression":"2 + 3 > 1"
							}`),
						RefID: "B",
						RelativeTimeRange: eval.RelativeTimeRange{
							From: eval.Duration(time.Duration(5) * time.Hour),
							To:   eval.Duration(time.Duration(3) * time.Hour),
						},
					},
				},
			}
			if tc.inputIntervalSeconds != nil {
				q.IntervalSeconds = tc.inputIntervalSeconds
			}
			err := ng.saveAlertDefinition(&q)
			switch {
			case tc.expectedError != nil:
				require.Error(t, err)
			default:
				require.NoError(t, err)
				assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
				assert.Equal(t, tc.expectedInterval, q.Result.IntervalSeconds)
				assert.Equal(t, int64(1), q.Result.Version)

			}
		})
	}
}
func TestCreatingConflictionAlertDefinition(t *testing.T) {
	t.Run("Should fail to create alert definition with conflicting org_id, title", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := saveAlertDefinitionCommand{
			OrgID:     1,
			Title:     "title",
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
								"datasource": "__expr__",
								"type":"math",
								"expression":"2 + 3 > 1"
							}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(time.Duration(5) * time.Hour),
						To:   eval.Duration(time.Duration(3) * time.Hour),
					},
				},
			},
		}

		err := ng.saveAlertDefinition(&q)
		require.NoError(t, err)

		err = ng.saveAlertDefinition(&q)
		require.Error(t, err)
		assert.True(t, ng.SQLStore.Dialect.IsUniqueConstraintViolation(err))
	})
}

func TestUpdatingAlertDefinition(t *testing.T) {
	t.Run("zero rows affected when updating unknown alert", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := updateAlertDefinitionCommand{
			UID:       "unknown",
			OrgID:     1,
			Title:     "something completely different",
			Condition: "A",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 2 > 1"
						}`),
					RefID: "A",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(time.Duration(5) * time.Hour),
						To:   eval.Duration(time.Duration(3) * time.Hour),
					},
				},
			},
		}

		err := ng.updateAlertDefinition(&q)
		require.NoError(t, err)
	})

	t.Run("updating existing alert", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		var initialInterval int64 = 120
		alertDefinition := createTestAlertDefinition(t, ng, initialInterval)
		created := alertDefinition.Updated

		var customInterval int64 = 30
		testCases := []struct {
			desc                    string
			inputOrgID              int64
			inputTitle              string
			inputInterval           *int64
			expectedError           error
			expectedIntervalSeconds int64
			expectedUpdated         time.Time
			expectedTitle           string
		}{
			{
				desc:                    "should not update previous interval if it's not provided",
				inputInterval:           nil,
				inputOrgID:              alertDefinition.OrgID,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: initialInterval,
				expectedUpdated:         time.Unix(1, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:                    "should update interval if it's provided",
				inputInterval:           &customInterval,
				inputOrgID:              alertDefinition.OrgID,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(2, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:                    "should not update organisation if it's provided",
				inputInterval:           &customInterval,
				inputOrgID:              0,
				inputTitle:              "something completely different",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(3, 0).UTC(),
				expectedTitle:           "something completely different",
			},
			{
				desc:          "should not update alert definition if the title it's too big",
				inputInterval: &customInterval,
				inputOrgID:    0,
				inputTitle:    getLongString(alertDefinitionMaxTitleLength + 1),
				expectedError: errors.New(""),
			},
			{
				desc:                    "should not update alert definition title if the title is empty",
				inputInterval:           &customInterval,
				inputOrgID:              0,
				inputTitle:              "",
				expectedIntervalSeconds: customInterval,
				expectedUpdated:         time.Unix(4, 0).UTC(),
				expectedTitle:           "something completely different",
			},
		}

		q := updateAlertDefinitionCommand{
			UID:       (*alertDefinition).UID,
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 3 > 1"
						}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(5 * time.Hour),
						To:   eval.Duration(3 * time.Hour),
					},
				},
			},
		}

		lastUpdated := created
		previousAlertDefinition := alertDefinition
		for _, tc := range testCases {
			t.Run(tc.desc, func(t *testing.T) {
				if tc.inputInterval != nil {
					q.IntervalSeconds = tc.inputInterval
				}
				if tc.inputOrgID != 0 {
					q.OrgID = tc.inputOrgID
				}
				q.Title = tc.inputTitle
				err := ng.updateAlertDefinition(&q)
				switch {
				case tc.expectedError != nil:
					require.Error(t, err)

					assert.Equal(t, previousAlertDefinition.Title, q.Result.Title)
					assert.Equal(t, previousAlertDefinition.Condition, q.Result.Condition)
					assert.Equal(t, len(previousAlertDefinition.Data), len(q.Result.Data))
					assert.Equal(t, previousAlertDefinition.IntervalSeconds, q.Result.IntervalSeconds)
					assert.Equal(t, previousAlertDefinition.Updated, q.Result.Updated)
					assert.Equal(t, previousAlertDefinition.Version, q.Result.Version)
					assert.Equal(t, previousAlertDefinition.OrgID, q.Result.OrgID)
					assert.Equal(t, previousAlertDefinition.UID, q.Result.UID)
				default:
					require.NoError(t, err)
					assert.Equal(t, previousAlertDefinition.ID, q.Result.ID)
					assert.Equal(t, previousAlertDefinition.UID, q.Result.UID)
					assert.True(t, q.Result.Updated.After(lastUpdated))
					assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
					assert.Equal(t, previousAlertDefinition.Version+1, q.Result.Version)

					assert.Equal(t, alertDefinition.OrgID, q.Result.OrgID)

					assert.Equal(t, "something completely different", q.Result.Title)
					assert.Equal(t, "B", q.Result.Condition)
					assert.Equal(t, 1, len(q.Result.Data))
					assert.Equal(t, tc.expectedUpdated, q.Result.Updated)
					assert.Equal(t, tc.expectedIntervalSeconds, q.Result.IntervalSeconds)
					assert.Equal(t, previousAlertDefinition.Version+1, q.Result.Version)
					assert.Equal(t, alertDefinition.OrgID, q.Result.OrgID)
					assert.Equal(t, alertDefinition.UID, q.Result.UID)

					previousAlertDefinition = q.Result
				}
			})

		}

	})
}

func TestUpdatingConflictingAlertDefinition(t *testing.T) {
	t.Run("should fail to update alert definition with reserved title", func(t *testing.T) {
		mockTimeNow()
		defer resetTimeNow()

		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		var initialInterval int64 = 120
		alertDef1 := createTestAlertDefinition(t, ng, initialInterval)
		alertDef2 := createTestAlertDefinition(t, ng, initialInterval)

		q := updateAlertDefinitionCommand{
			UID:       (*alertDef2).UID,
			Title:     alertDef1.Title,
			Condition: "B",
			Data: []eval.AlertQuery{
				{
					Model: json.RawMessage(`{
							"datasource": "__expr__",
							"type":"math",
							"expression":"2 + 3 > 1"
						}`),
					RefID: "B",
					RelativeTimeRange: eval.RelativeTimeRange{
						From: eval.Duration(5 * time.Hour),
						To:   eval.Duration(3 * time.Hour),
					},
				},
			},
		}

		err := ng.updateAlertDefinition(&q)
		require.Error(t, err)
		assert.True(t, ng.SQLStore.Dialect.IsUniqueConstraintViolation(err))
	})
}

func TestDeletingAlertDefinition(t *testing.T) {
	t.Run("zero rows affected when deleting unknown alert", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		q := deleteAlertDefinitionByUIDCommand{
			UID:   "unknown",
			OrgID: 1,
		}

		err := ng.deleteAlertDefinitionByUID(&q)
		require.NoError(t, err)
	})

	t.Run("deleting successfully existing alert", func(t *testing.T) {
		ng := setupTestEnv(t)
		t.Cleanup(registry.ClearOverrides)

		alertDefinition := createTestAlertDefinition(t, ng, 60)

		q := deleteAlertDefinitionByUIDCommand{
			UID:   (*alertDefinition).UID,
			OrgID: 1,
		}

		// save an instance for the definition
		saveCmd := &saveAlertInstanceCommand{
			DefinitionOrgID: alertDefinition.OrgID,
			DefinitionUID:   alertDefinition.UID,
			State:           InstanceStateFiring,
			Labels:          InstanceLabels{"test": "testValue"},
		}
		err := ng.saveAlertInstance(saveCmd)
		require.NoError(t, err)
		listCommand := &listAlertInstancesQuery{
			DefinitionOrgID: alertDefinition.OrgID,
			DefinitionUID:   alertDefinition.UID,
		}
		err = ng.listAlertInstances(listCommand)
		require.NoError(t, err)
		require.Len(t, listCommand.Result, 1)

		err = ng.deleteAlertDefinitionByUID(&q)
		require.NoError(t, err)

		// assert that alert instance is deleted
		err = ng.listAlertInstances(listCommand)
		require.NoError(t, err)

		require.Len(t, listCommand.Result, 0)
	})
}

func getLongString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = 'a'
	}
	return string(b)
}
<|endoftext|>


<commit_before>package infrastructure

import (
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/satori/go.uuid"
	"github.com/zeph/goad/version"
)

// Infrastructure manages the resource creation and updates necessary to use
// Goad.
type Infrastructure struct {
	config   *aws.Config
	queueURL string
	regions  []string
}

// New creates the required infrastructure to run the load tests in Lambda
// functions.
func New(regions []string, config *aws.Config) (*Infrastructure, error) {
	infra := &Infrastructure{config: config, regions: regions}
	if err := infra.setup(); err != nil {
		return nil, err
	}
	return infra, nil
}

// QueueURL returns the URL of the SQS queue to use for the load test session.
func (infra *Infrastructure) QueueURL() string {
	return infra.queueURL
}

// Clean removes any AWS resources that cannot be reused for a subsequent test.
func (infra *Infrastructure) Clean() {
	infra.removeSQSQueue()
}

func (infra *Infrastructure) setup() error {
	roleArn, err := infra.createIAMLambdaRole("goad-lambda-role")
	if err != nil {
		return err
	}
	zip, err := Asset("data/lambda.zip")
	if err != nil {
		return err
	}

	for _, region := range infra.regions {
		err = infra.createOrUpdateLambdaFunction(region, roleArn, zip)
		if err != nil {
			return err
		}
	}
	queueURL, err := infra.createSQSQueue()
	if err != nil {
		return err
	}
	infra.queueURL = queueURL
	return nil
}

func (infra *Infrastructure) createOrUpdateLambdaFunction(region, roleArn string, payload []byte) error {
	config := aws.NewConfig().WithRegion(region)
	svc := lambda.New(session.New(), config)

	exists, err := lambdaExists(svc)

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

	if exists {
		aliasExists, err := lambdaAliasExists(svc)
		if err != nil || aliasExists {
			return infra.updateLambdaFunction(svc, roleArn, payload)
		}
	}

	return infra.createLambdaFunction(svc, roleArn, payload)
}

func (infra *Infrastructure) createLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.CreateFunction(&lambda.CreateFunctionInput{
		Code: &lambda.FunctionCode{
			ZipFile: payload,
		},
		FunctionName: aws.String("goad"),
		Handler:      aws.String("index.handler"),
		Role:         aws.String(roleArn),
		Runtime:      aws.String("nodejs"),
		MemorySize:   aws.Int64(1536),
		Publish:      aws.Bool(true),
		Timeout:      aws.Int64(300),
	})
	if err != nil {
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func (infra *Infrastructure) updateLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		ZipFile:      payload,
		FunctionName: aws.String("goad"),
		Publish:      aws.Bool(true),
	})
	if err != nil {
		return err
	}
	return updateLambdaAlias(svc, function.Version)
}

func lambdaExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetFunction(&lambda.GetFunctionInput{
		FunctionName: aws.String("goad"),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func createLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func updateLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func lambdaAliasExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetAlias(&lambda.GetAliasInput{
		FunctionName: aws.String("goad"),
		Name:         aws.String(version.LambdaVersion()),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func (infra *Infrastructure) createIAMLambdaRole(roleName string) (arn string, err error) {
	svc := iam.New(session.New(), infra.config)

	resp, err := svc.GetRole(&iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				res, err := svc.CreateRole(&iam.CreateRoleInput{
					AssumeRolePolicyDocument: aws.String(`{
        	          "Version": "2012-10-17",
        	          "Statement": {
        	            "Effect": "Allow",
        	            "Principal": {"Service": "lambda.amazonaws.com"},
        	            "Action": "sts:AssumeRole"
        	          }
            	    }`),
					RoleName: aws.String(roleName),
					Path:     aws.String("/"),
				})
				if err != nil {
					return "", err
				}
				if err := infra.createIAMLambdaRolePolicy(*res.Role.RoleName); err != nil {
					return "", err
				}
				return *res.Role.Arn, nil
			}
		}
		return "", err
	}

	return *resp.Role.Arn, nil
}

func (infra *Infrastructure) createIAMLambdaRolePolicy(roleName string) error {
	svc := iam.New(session.New(), infra.config)

	_, err := svc.PutRolePolicy(&iam.PutRolePolicyInput{
		PolicyDocument: aws.String(`{
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "sqs:SendMessage"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:sqs:*:*:goad-*"
		  	},
			{
              "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:logs:*:*:*"
	        }
          ]
        }`),
		PolicyName: aws.String("goad-lambda-role-policy"),
		RoleName:   aws.String(roleName),
	})
	return err
}

func (infra *Infrastructure) createSQSQueue() (url string, err error) {
	svc := sqs.New(session.New(), infra.config)

	resp, err := svc.CreateQueue(&sqs.CreateQueueInput{
		QueueName: aws.String("goad-" + uuid.NewV4().String()),
	})

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

	return *resp.QueueUrl, nil
}

func (infra *Infrastructure) removeSQSQueue() {
	svc := sqs.New(session.New(), infra.config)

	svc.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(infra.queueURL),
	})
}
<commit_msg>and, apparently nodejs(!4.3) was deprecated<commit_after>package infrastructure

import (
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/satori/go.uuid"
	"github.com/zeph/goad/version"
)

// Infrastructure manages the resource creation and updates necessary to use
// Goad.
type Infrastructure struct {
	config   *aws.Config
	queueURL string
	regions  []string
}

// New creates the required infrastructure to run the load tests in Lambda
// functions.
func New(regions []string, config *aws.Config) (*Infrastructure, error) {
	infra := &Infrastructure{config: config, regions: regions}
	if err := infra.setup(); err != nil {
		return nil, err
	}
	return infra, nil
}

// QueueURL returns the URL of the SQS queue to use for the load test session.
func (infra *Infrastructure) QueueURL() string {
	return infra.queueURL
}

// Clean removes any AWS resources that cannot be reused for a subsequent test.
func (infra *Infrastructure) Clean() {
	infra.removeSQSQueue()
}

func (infra *Infrastructure) setup() error {
	roleArn, err := infra.createIAMLambdaRole("goad-lambda-role")
	if err != nil {
		return err
	}
	zip, err := Asset("data/lambda.zip")
	if err != nil {
		return err
	}

	for _, region := range infra.regions {
		err = infra.createOrUpdateLambdaFunction(region, roleArn, zip)
		if err != nil {
			return err
		}
	}
	queueURL, err := infra.createSQSQueue()
	if err != nil {
		return err
	}
	infra.queueURL = queueURL
	return nil
}

func (infra *Infrastructure) createOrUpdateLambdaFunction(region, roleArn string, payload []byte) error {
	config := aws.NewConfig().WithRegion(region)
	svc := lambda.New(session.New(), config)

	exists, err := lambdaExists(svc)

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

	if exists {
		aliasExists, err := lambdaAliasExists(svc)
		if err != nil || aliasExists {
			return infra.updateLambdaFunction(svc, roleArn, payload)
		}
	}

	return infra.createLambdaFunction(svc, roleArn, payload)
}

func (infra *Infrastructure) createLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.CreateFunction(&lambda.CreateFunctionInput{
		Code: &lambda.FunctionCode{
			ZipFile: payload,
		},
		FunctionName: aws.String("goad"),
		Handler:      aws.String("index.handler"),
		Role:         aws.String(roleArn),
		Runtime:      aws.String("nodejs4.3"),
		MemorySize:   aws.Int64(1536),
		Publish:      aws.Bool(true),
		Timeout:      aws.Int64(300),
	})
	if err != nil {
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func (infra *Infrastructure) updateLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		ZipFile:      payload,
		FunctionName: aws.String("goad"),
		Publish:      aws.Bool(true),
	})
	if err != nil {
		return err
	}
	return updateLambdaAlias(svc, function.Version)
}

func lambdaExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetFunction(&lambda.GetFunctionInput{
		FunctionName: aws.String("goad"),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func createLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func updateLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func lambdaAliasExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetAlias(&lambda.GetAliasInput{
		FunctionName: aws.String("goad"),
		Name:         aws.String(version.LambdaVersion()),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func (infra *Infrastructure) createIAMLambdaRole(roleName string) (arn string, err error) {
	svc := iam.New(session.New(), infra.config)

	resp, err := svc.GetRole(&iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				res, err := svc.CreateRole(&iam.CreateRoleInput{
					AssumeRolePolicyDocument: aws.String(`{
        	          "Version": "2012-10-17",
        	          "Statement": {
        	            "Effect": "Allow",
        	            "Principal": {"Service": "lambda.amazonaws.com"},
        	            "Action": "sts:AssumeRole"
        	          }
            	    }`),
					RoleName: aws.String(roleName),
					Path:     aws.String("/"),
				})
				if err != nil {
					return "", err
				}
				if err := infra.createIAMLambdaRolePolicy(*res.Role.RoleName); err != nil {
					return "", err
				}
				return *res.Role.Arn, nil
			}
		}
		return "", err
	}

	return *resp.Role.Arn, nil
}

func (infra *Infrastructure) createIAMLambdaRolePolicy(roleName string) error {
	svc := iam.New(session.New(), infra.config)

	_, err := svc.PutRolePolicy(&iam.PutRolePolicyInput{
		PolicyDocument: aws.String(`{
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "sqs:SendMessage"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:sqs:*:*:goad-*"
		  	},
			{
              "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:logs:*:*:*"
	        }
          ]
        }`),
		PolicyName: aws.String("goad-lambda-role-policy"),
		RoleName:   aws.String(roleName),
	})
	return err
}

func (infra *Infrastructure) createSQSQueue() (url string, err error) {
	svc := sqs.New(session.New(), infra.config)

	resp, err := svc.CreateQueue(&sqs.CreateQueueInput{
		QueueName: aws.String("goad-" + uuid.NewV4().String()),
	})

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

	return *resp.QueueUrl, nil
}

func (infra *Infrastructure) removeSQSQueue() {
	svc := sqs.New(session.New(), infra.config)

	svc.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(infra.queueURL),
	})
}
<|endoftext|>


<commit_before>package gobeacon

import (
	"errors"
	"net/url"
	"packetbeat/common"
	"packetbeat/logp"
	"reflect"
	"testing"

	"go-beacon-orig/go/codec"

	"github.com/gdamore/mangos/protocol/req"
	"github.com/gdamore/mangos/transport/all"
	"github.com/stretchr/testify/assert"
)

func sendMessage(url string, message *[]byte) error {

	requestSocket, err := req.NewSocket()
	if err != nil {
		return err
	}
	defer requestSocket.Close()
	all.AddTransports(requestSocket)

	if err = requestSocket.Dial(url); err != nil {
		return err
	}

	if err = requestSocket.Send(*message); err != nil {
		return err
	}

	var clientMsg []byte

	if clientMsg, err = requestSocket.Recv(); err != nil {
		return err
	}

	if string(clientMsg) != "OK" {
		return errors.New("Response not OK, requeued")
	}

	return nil

}

func encodeMessage(query url.Values) ([]byte, error) {
	var b []byte
	var mh codec.MsgpackHandle
	mh.MapType = reflect.TypeOf(map[string][]string(nil))

	enc := codec.NewEncoderBytes(&b, &mh)
	err := enc.Encode(query)
	if err != nil {
		return nil, err
	}
	return b, nil

}

func TestGoBeacon(t *testing.T) {
	if testing.Verbose() {
		logp.LogInit(logp.LOG_DEBUG, "", false, true, []string{"gobeacon"})
	}

	events := make(chan common.MapStr)
	server := new(GoBeacon)

	server.Config = Config{
		ListenAddr: "tcp://127.0.0.1:9000",
		Tracker:    "boomerang",
	}
	err := server.Init(true, events)
	assert.Nil(t, err)

	ready := make(chan bool)

	go func() {
		ready <- true
		err := server.Run()
		assert.Nil(t, err, "Error: %v", err)
	}()

	// make sure the goroutine runs first
	<-ready

	// send a message
	v := url.Values{}
	v.Set("nt_nav_type", "1")
	v.Add("nt_domcontloaded_st", "1426674694041")
	v.Add("nt_domcontloaded_end", "1426674694041")
	v.Add("u", "http://localhost:8080/static/")
	v.Add("r", "")
	v.Add("nt_domcomp", "1426674694056")
	v.Add("nt_dns_st", "1426674694014")
	v.Add("nt_dns_end", "1426675597491")
	v.Add("rt.bstart", "1426674694048")
	v.Add("rt.end", "1426674694057")
	v.Add("nt_con_st", "1426675597491")
	v.Add("nt_con_end", "1426675597491")

	obj, err := encodeMessage(v)
	err = sendMessage(server.Config.ListenAddr, &obj)
	assert.Nil(t, err, "Error: %v", err)

	server.Stop()
}
<commit_msg>Import the right package<commit_after>package gobeacon

import (
	"errors"
	"net/url"
	"packetbeat/common"
	"packetbeat/logp"
	"reflect"
	"testing"

	"github.com/gdamore/mangos/protocol/req"
	"github.com/gdamore/mangos/transport/all"
	"github.com/stretchr/testify/assert"
	"github.com/ugorji/go/codec"
)

func sendMessage(url string, message *[]byte) error {

	requestSocket, err := req.NewSocket()
	if err != nil {
		return err
	}
	defer requestSocket.Close()
	all.AddTransports(requestSocket)

	if err = requestSocket.Dial(url); err != nil {
		return err
	}

	if err = requestSocket.Send(*message); err != nil {
		return err
	}

	var clientMsg []byte

	if clientMsg, err = requestSocket.Recv(); err != nil {
		return err
	}

	if string(clientMsg) != "OK" {
		return errors.New("Response not OK, requeued")
	}

	return nil

}

func encodeMessage(query url.Values) ([]byte, error) {
	var b []byte
	var mh codec.MsgpackHandle
	mh.MapType = reflect.TypeOf(map[string][]string(nil))

	enc := codec.NewEncoderBytes(&b, &mh)
	err := enc.Encode(query)
	if err != nil {
		return nil, err
	}
	return b, nil

}

func TestGoBeacon(t *testing.T) {
	if testing.Verbose() {
		logp.LogInit(logp.LOG_DEBUG, "", false, true, []string{"gobeacon"})
	}

	events := make(chan common.MapStr)
	server := new(GoBeacon)

	server.Config = Config{
		ListenAddr: "tcp://127.0.0.1:9000",
		Tracker:    "boomerang",
	}
	err := server.Init(true, events)
	assert.Nil(t, err)

	ready := make(chan bool)

	go func() {
		ready <- true
		err := server.Run()
		assert.Nil(t, err, "Error: %v", err)
	}()

	// make sure the goroutine runs first
	<-ready

	// send a message
	v := url.Values{}
	v.Set("nt_nav_type", "1")
	v.Add("nt_domcontloaded_st", "1426674694041")
	v.Add("nt_domcontloaded_end", "1426674694041")
	v.Add("u", "http://localhost:8080/static/")
	v.Add("r", "")
	v.Add("nt_domcomp", "1426674694056")
	v.Add("nt_dns_st", "1426674694014")
	v.Add("nt_dns_end", "1426675597491")
	v.Add("rt.bstart", "1426674694048")
	v.Add("rt.end", "1426674694057")
	v.Add("nt_con_st", "1426675597491")
	v.Add("nt_con_end", "1426675597491")

	obj, err := encodeMessage(v)
	err = sendMessage(server.Config.ListenAddr, &obj)
	assert.Nil(t, err, "Error: %v", err)

	server.Stop()
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.0"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = "dev"

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<commit_msg>v0.9.0<commit_after>package terraform

import (
	"fmt"

	"github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
const Version = "0.9.0"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
const VersionPrerelease = ""

// SemVersion is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVersion = version.Must(version.NewVersion(Version))

// VersionHeader is the header name used to send the current terraform version
// in http requests.
const VersionHeader = "Terraform-Version"

func VersionString() string {
	if VersionPrerelease != "" {
		return fmt.Sprintf("%s-%s", Version, VersionPrerelease)
	}
	return Version
}
<|endoftext|>


<commit_before>package queue

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

	"github.com/garyburd/redigo/redis"
	"github.com/pborman/uuid"
	log "github.com/sirupsen/logrus"
)

func ListenAndServe(addr string) error {
	http.HandleFunc("/create", createHandler)
	http.HandleFunc("/delete", deleteHandler)
	http.HandleFunc("/query", queryHandler)
	return http.ListenAndServe(addr, nil)
}

type createRequest struct {
	Topic string `json:"topic"`
	// Delay is the number of seconds that should elapse before the task execute
	Delay    int64  `json:"delay"`
	Retry    int    `json:"retry"`
	Callback string `json:"callback"`
	Content  string `json:"content"`
}

type createResponse struct {
	ID string `json:"id"`
}

func createHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request createRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	task := &Task{
		ID:          uuid.New(),
		Topic:       request.Topic,
		ExecuteTime: time.Now().Unix() + request.Delay,
		MaxRetry:    request.Retry,
		Callback:    request.Callback,
		Content:     request.Content,
		CreatTime:   time.Now().Unix(),
	}
	err = createTask(task)
	if err != nil {
		log.WithError(err).Error("create task fail")
		w.WriteHeader(500)
		return
	}
	response := createResponse{ID: task.ID}
	respData, err := json.Marshal(response)
	if err != nil {
		log.WithError(err).Error("json marshal fail")
		w.WriteHeader(500)
		return
	}
	w.Write(respData)
}

type deleteRequest struct {
	ID string `json:"id"`
}

func deleteHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request deleteRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	if request.ID == "" {
		w.WriteHeader(400)
		return
	}
	err = deleteTask(request.ID)
	if err != nil {
		log.WithError(err).Error("delete task fail")
		w.WriteHeader(500)
		return
	}
	w.WriteHeader(200)
}

type queryRequest struct {
	ID string `json:"id"`
}

type queryResponse struct {
	ID          string `json:"id"`
	Topic       string `json:"topic"`
	ExecuteTime int64  `json:"execute_time"`
	MaxRetry    int    `json:"max_retry"`
	HasRetry    int    `json:"has_retry"`
	Callback    string `json:"callback"`
	Content     string `json:"content"`
	CreatTime   int64  `json:"creat_time"`
}

func queryHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request queryRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	if request.ID == "" {
		w.WriteHeader(400)
		return
	}
	task, err := getTask(request.ID)
	if err != nil {
		if err == redis.ErrNil {
			w.WriteHeader(404)
			return
		}
		log.WithError(err).Error("get task fail")
		w.WriteHeader(500)
		return
	}
	response := queryResponse{
		ID:          task.ID,
		Topic:       task.Topic,
		ExecuteTime: task.ExecuteTime,
		MaxRetry:    task.MaxRetry,
		HasRetry:    task.HasRetry,
		Callback:    task.Callback,
		Content:     task.Content,
		CreatTime:   task.CreatTime,
	}
	respData, err := json.Marshal(response)
	if err != nil {
		log.WithError(err).Error("json marshal fail")
		w.WriteHeader(500)
		return
	}
	w.Write(respData)
}
<commit_msg>feat: use graceful http<commit_after>package queue

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

	"github.com/facebookgo/grace/gracehttp"
	"github.com/garyburd/redigo/redis"
	"github.com/pborman/uuid"
	log "github.com/sirupsen/logrus"
)

func ListenAndServe(addr string) error {
	mux := http.NewServeMux()
	mux.HandleFunc("/create", createHandler)
	mux.HandleFunc("/delete", deleteHandler)
	mux.HandleFunc("/query", queryHandler)
	server := http.Server{Addr: addr, Handler: mux}
	return gracehttp.Serve(&server)
}

type createRequest struct {
	Topic string `json:"topic"`
	// Delay is the number of seconds that should elapse before the task execute
	Delay    int64  `json:"delay"`
	Retry    int    `json:"retry"`
	Callback string `json:"callback"`
	Content  string `json:"content"`
}

type createResponse struct {
	ID string `json:"id"`
}

func createHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request createRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	task := &Task{
		ID:          uuid.New(),
		Topic:       request.Topic,
		ExecuteTime: time.Now().Unix() + request.Delay,
		MaxRetry:    request.Retry,
		Callback:    request.Callback,
		Content:     request.Content,
		CreatTime:   time.Now().Unix(),
	}
	err = createTask(task)
	if err != nil {
		log.WithError(err).Error("create task fail")
		w.WriteHeader(500)
		return
	}
	response := createResponse{ID: task.ID}
	respData, err := json.Marshal(response)
	if err != nil {
		log.WithError(err).Error("json marshal fail")
		w.WriteHeader(500)
		return
	}
	w.Write(respData)
}

type deleteRequest struct {
	ID string `json:"id"`
}

func deleteHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request deleteRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	if request.ID == "" {
		w.WriteHeader(400)
		return
	}
	err = deleteTask(request.ID)
	if err != nil {
		log.WithError(err).Error("delete task fail")
		w.WriteHeader(500)
		return
	}
	w.WriteHeader(200)
}

type queryRequest struct {
	ID string `json:"id"`
}

type queryResponse struct {
	ID          string `json:"id"`
	Topic       string `json:"topic"`
	ExecuteTime int64  `json:"execute_time"`
	MaxRetry    int    `json:"max_retry"`
	HasRetry    int    `json:"has_retry"`
	Callback    string `json:"callback"`
	Content     string `json:"content"`
	CreatTime   int64  `json:"creat_time"`
}

func queryHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		w.WriteHeader(400)
		return
	}
	var request queryRequest
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.WithError(err).Error("io read from frontend fail")
		w.WriteHeader(500)
		return
	}
	err = json.Unmarshal(data, &request)
	if err != nil {
		log.WithError(err).Error("json unmarshal fail")
		w.WriteHeader(400)
		return
	}
	if request.ID == "" {
		w.WriteHeader(400)
		return
	}
	task, err := getTask(request.ID)
	if err != nil {
		if err == redis.ErrNil {
			w.WriteHeader(404)
			return
		}
		log.WithError(err).Error("get task fail")
		w.WriteHeader(500)
		return
	}
	response := queryResponse{
		ID:          task.ID,
		Topic:       task.Topic,
		ExecuteTime: task.ExecuteTime,
		MaxRetry:    task.MaxRetry,
		HasRetry:    task.HasRetry,
		Callback:    task.Callback,
		Content:     task.Content,
		CreatTime:   task.CreatTime,
	}
	respData, err := json.Marshal(response)
	if err != nil {
		log.WithError(err).Error("json marshal fail")
		w.WriteHeader(500)
		return
	}
	w.Write(respData)
}
<|endoftext|>


<commit_before>package quickfix

import "honnef.co/go/tools/analysis/lint"

var Docs = lint.Markdownify(map[string]*lint.Documentation{
	"QF1000": {
		Title: "Use byte-specific indexing function",
		Since: "Unreleased",
	},
	"QF1001": {
		Title: "Apply De Morgan's law",
		Since: "Unreleased",
	},
	"QF1002": {
		Title: "Convert untagged switch to tagged switch",
		Text: `An untagged switch that compares a single variable against a series of values can be replaced with a tagged switch.

Before:

	switch {
	case x == 1 || x == 2, x == 3:
		...
	case x == 4:
		...
	default:
		...
	}

After:

	switch x {
	case 1, 2, 3:
		...
	case 4:
		...
	default:
		...
	}
`,
		Since: "Unreleased",
	},
	"QF1003": {
		Title: "Convert if/else-if chain to tagged switch",
		Text: `A series of if/else-if checks comparing the same variable against values can be replaced with a tagged switch.

Before:

	if x == 1 || x == 2 {
		...
	} else if x == 3 {
		...
	} else {
		...
	}

After:

	switch x {
	case 1, 2:
		...
	case 3:
		...
	default:
		...
	}`,
		Since: "Unreleased",
	},
	"QF1004": {
		Since: "Unreleased",
	},
	"QF1005": {
		Title: "Expand call to `math.Pow`",
		Text: `Some uses of \'math.Pow\' can be simplified to basic multiplication.

Before:

	math.Pow(x, 2)

After:

	x * x`,
		Since: "Unreleased",
	},

	"QF1006": {
		Title: "Lift if+break into loop condition",
		Text: `Before:

	for {
		if done {
			break
		}
		...
	}

After:

	for !done {
		...
	}`,
		Since: "Unreleased",
	},

	"QF1007": {
		Title: "Merge conditional assignment into variable declaration",
		Text: `Before:

	x := false
	if someCondition {
		x = true
	}

After:

	x := someCondition
`,
		Since: "Unreleased",
	},

	"QF1008": {
		Title: "Omit embedded fields from selector expression",
		Since: "Unreleased",
	},

	"QF1009": {
		Title: `Use time.Time.Equal instead of \'==\' operator`,
		Since: "Unreleased",
	},

	"QF1010": {
		Title: "Convert slice of bytes to string when printing it",
		Since: "Unreleased",
	},
})
<commit_msg>quickfix: add severities to all checks<commit_after>package quickfix

import "honnef.co/go/tools/analysis/lint"

var Docs = lint.Markdownify(map[string]*lint.Documentation{
	"QF1000": {
		Title:    "Use byte-specific indexing function",
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1001": {
		Title:    "Apply De Morgan's law",
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1002": {
		Title: "Convert untagged switch to tagged switch",
		Text: `An untagged switch that compares a single variable against a series of values can be replaced with a tagged switch.

Before:

	switch {
	case x == 1 || x == 2, x == 3:
		...
	case x == 4:
		...
	default:
		...
	}

After:

	switch x {
	case 1, 2, 3:
		...
	case 4:
		...
	default:
		...
	}
`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1003": {
		Title: "Convert if/else-if chain to tagged switch",
		Text: `A series of if/else-if checks comparing the same variable against values can be replaced with a tagged switch.

Before:

	if x == 1 || x == 2 {
		...
	} else if x == 3 {
		...
	} else {
		...
	}

After:

	switch x {
	case 1, 2:
		...
	case 3:
		...
	default:
		...
	}`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1004": {
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1005": {
		Title: "Expand call to `math.Pow`",
		Text: `Some uses of \'math.Pow\' can be simplified to basic multiplication.

Before:

	math.Pow(x, 2)

After:

	x * x`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1006": {
		Title: "Lift if+break into loop condition",
		Text: `Before:

	for {
		if done {
			break
		}
		...
	}

After:

	for !done {
		...
	}`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1007": {
		Title: "Merge conditional assignment into variable declaration",
		Text: `Before:

	x := false
	if someCondition {
		x = true
	}

After:

	x := someCondition
`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1008": {
		Title:    "Omit embedded fields from selector expression",
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1009": {
		Title:    `Use time.Time.Equal instead of \'==\' operator`,
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},

	"QF1010": {
		Title:    "Convert slice of bytes to string when printing it",
		Since:    "Unreleased",
		Severity: lint.SeverityHint,
	},
})
<|endoftext|>


<commit_before>package kite

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/tools/amqputil"
	"koding/tools/dnode"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type Kite struct {
	Name     string
	Handlers map[string]Handler
}

type Handler struct {
	Concurrent bool
	Callback   func(args *dnode.Partial, session *Session) (interface{}, error)
}

func New(name string) *Kite {
	return &Kite{
		Name:     name,
		Handlers: make(map[string]Handler),
	}
}

func (k *Kite) Handle(method string, concurrent bool, callback func(args *dnode.Partial, session *Session) (interface{}, error)) {
	k.Handlers[method] = Handler{concurrent, callback}
}

func (k *Kite) Run() {
	changeClientsGauge := lifecycle.CreateClientsGauge()
	log.RunGaugesLoop()

	routeMap := make(map[string](chan<- []byte))
	defer func() {
		for _, channel := range routeMap {
			close(channel)
		}
	}()

	timeoutChannel := make(chan string)

	sigtermChannel := make(chan os.Signal)
	signal.Notify(sigtermChannel, syscall.SIGTERM)

	consumeConn := amqputil.CreateConnection("kite-" + k.Name)
	defer consumeConn.Close()

	publishConn := amqputil.CreateConnection("kite-" + k.Name)
	defer publishConn.Close()

	publishChannel := amqputil.CreateChannel(publishConn)
	defer publishChannel.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	amqputil.DeclarePresenceExchange(consumeChannel, "services-presence", "kite", "kite-"+k.Name, "kite-"+k.Name)
	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "fanout", "kite-"+k.Name, "", true)

	for {
		select {
		case message, ok := <-stream:
			if !ok {
				return
			}

			switch message.RoutingKey {
			case "auth.join":
				var client struct {
					Username   string
					RoutingKey string
				}
				err := json.Unmarshal(message.Body, &client)
				if err != nil || client.Username == "" || client.RoutingKey == "" {
					log.Err("Invalid auth.join message.", message.Body)
					continue
				}

				if _, found := routeMap[client.RoutingKey]; found {
					log.Warn("Duplicate auth.join for same routing key.")
					continue
				}
				channel := make(chan []byte, 1024)
				routeMap[client.RoutingKey] = channel

				go func() {
					defer log.RecoverAndLog()

					changeClientsGauge(1)
					log.Debug("Client connected: " + client.Username)
					defer func() {
						changeClientsGauge(-1)
						log.Debug("Client disconnected: " + client.Username)
					}()

					session := NewSession(client.Username)
					defer session.Close()

					d := dnode.New()
					defer d.Close()
					d.OnRootMethod = func(method string, args *dnode.Partial) {
						defer log.RecoverAndLog()

						if method == "ping" {
							d.Send("pong")
							return
						}

						var partials []*dnode.Partial
						err := args.Unmarshal(&partials)
						if err != nil {
							panic(err)
						}

						var options struct {
							WithArgs *dnode.Partial
						}
						err = partials[0].Unmarshal(&options)
						if err != nil {
							panic(err)
						}
						var resultCallback dnode.Callback
						err = partials[1].Unmarshal(&resultCallback)
						if err != nil {
							panic(err)
						}

						handler, found := k.Handlers[method]
						if !found {
							resultCallback(fmt.Sprintf("Method '%v' not known.", method), nil)
							return
						}

						execHandler := func() {
							result, err := handler.Callback(options.WithArgs, session)
							if b, ok := result.([]byte); ok {
								result = string(b)
							}

							if err != nil {
								resultCallback(err.Error(), result)
								return
							}

							resultCallback(nil, result)
						}

						if handler.Concurrent {
							go func() {
								defer log.RecoverAndLog()
								execHandler()
							}()
							return
						}

						execHandler()
					}

					go func() {
						defer log.RecoverAndLog()
						for data := range d.SendChan {
							log.Debug("Write", client.RoutingKey, data)
							err := publishChannel.Publish("broker", client.RoutingKey, false, false, amqp.Publishing{Body: data})
							if err != nil {
								log.LogError(err, 0)
							}
						}
					}()

					d.Send("ready", "kite-"+k.Name)

					for {
						select {
						case message, ok := <-channel:
							if !ok {
								return
							}
							log.Debug("Read", client.RoutingKey, message)
							d.ProcessMessage(message)
						case <-time.After(24 * time.Hour):
							timeoutChannel <- client.RoutingKey
						}
					}
				}()

			case "auth.leave":
				var client struct {
					RoutingKey string
				}
				err := json.Unmarshal(message.Body, &client)
				if err != nil || client.RoutingKey == "" {
					log.Err("Invalid auth.leave message.", message.Body)
					continue
				}

				channel, found := routeMap[client.RoutingKey]
				if found {
					close(channel)
					delete(routeMap, client.RoutingKey)
				}

			default:
				channel, found := routeMap[message.RoutingKey]
				if found {
					select {
					case channel <- message.Body:
						// successful
					default:
						close(channel)
						delete(routeMap, message.RoutingKey)
						log.Warn("Dropped client because of message buffer overflow.")
					}
				}
			}

		case routingKey := <-timeoutChannel:
			channel, found := routeMap[routingKey]
			if found {
				close(channel)
				delete(routeMap, routingKey)
				log.Warn("Dropped client because of fallback session timeout.")
			}

		case <-sigtermChannel:
			log.Info("Received TERM signal. Beginning shutdown...")
			lifecycle.BeginShutdown()
			consumeChannel.Close()
		}
	}
}

type Session struct {
	Username     string
	Alive        bool
	onDisconnect []func()
}

func NewSession(username string) *Session {
	return &Session{
		Username: username,
		Alive:    true,
	}
}

func (session *Session) OnDisconnect(f func()) {
	session.onDisconnect = append(session.onDisconnect, f)
}

func (session *Session) Close() {
	session.Alive = false
	for _, f := range session.onDisconnect {
		f()
	}
	session.onDisconnect = nil
}
<commit_msg>kite.go: work toward auth.who / kite.who<commit_after>package kite

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/tools/amqputil"
	"koding/tools/dnode"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type Kite struct {
	Name          string
	Handlers      map[string]Handler
	LoadBalancing bool
}

type Handler struct {
	Concurrent bool
	Callback   func(args *dnode.Partial, session *Session) (interface{}, error)
}

func New(name string) *Kite {
	return &Kite{
		Name:     name,
		Handlers: make(map[string]Handler),
	}
}

func (k *Kite) LoadBalance() {
	k.LoadBalancing = true
}

func (k *Kite) Handle(method string, concurrent bool, callback func(args *dnode.Partial, session *Session) (interface{}, error)) {
	k.Handlers[method] = Handler{concurrent, callback}
}

func (k *Kite) Run() {
	changeClientsGauge := lifecycle.CreateClientsGauge()
	log.RunGaugesLoop()

	routeMap := make(map[string](chan<- []byte))
	defer func() {
		for _, channel := range routeMap {
			close(channel)
		}
	}()

	timeoutChannel := make(chan string)

	sigtermChannel := make(chan os.Signal)
	signal.Notify(sigtermChannel, syscall.SIGTERM)

	consumeConn := amqputil.CreateConnection("kite-" + k.Name)
	defer consumeConn.Close()

	publishConn := amqputil.CreateConnection("kite-" + k.Name)
	defer publishConn.Close()

	publishChannel := amqputil.CreateChannel(publishConn)
	defer publishChannel.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	amqputil.DeclarePresenceExchange(consumeChannel, "services-presence", "kite", "kite-"+k.Name, "kite-"+k.Name, k.LoadBalancing)
	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "fanout", "kite-"+k.Name, "", true)

	for {
		select {
		case message, ok := <-stream:
			if !ok {
				return
			}

			switch message.RoutingKey {
			case "auth.join":
				var client struct {
					Username   string
					RoutingKey string
				}
				err := json.Unmarshal(message.Body, &client)
				if err != nil || client.Username == "" || client.RoutingKey == "" {
					log.Err("Invalid auth.join message.", message.Body)
					continue
				}

				if _, found := routeMap[client.RoutingKey]; found {
					log.Warn("Duplicate auth.join for same routing key.")
					continue
				}
				channel := make(chan []byte, 1024)
				routeMap[client.RoutingKey] = channel

				go func() {
					defer log.RecoverAndLog()

					changeClientsGauge(1)
					log.Debug("Client connected: " + client.Username)
					defer func() {
						changeClientsGauge(-1)
						log.Debug("Client disconnected: " + client.Username)
					}()

					session := NewSession(client.Username)
					defer session.Close()

					d := dnode.New()
					defer d.Close()
					d.OnRootMethod = func(method string, args *dnode.Partial) {
						defer log.RecoverAndLog()

						if method == "ping" {
							d.Send("pong")
							return
						}

						var partials []*dnode.Partial
						err := args.Unmarshal(&partials)
						if err != nil {
							panic(err)
						}

						var options struct {
							WithArgs *dnode.Partial
						}
						err = partials[0].Unmarshal(&options)
						if err != nil {
							panic(err)
						}
						var resultCallback dnode.Callback
						err = partials[1].Unmarshal(&resultCallback)
						if err != nil {
							panic(err)
						}

						handler, found := k.Handlers[method]
						if !found {
							resultCallback(fmt.Sprintf("Method '%v' not known.", method), nil)
							return
						}

						execHandler := func() {
							result, err := handler.Callback(options.WithArgs, session)
							if b, ok := result.([]byte); ok {
								result = string(b)
							}

							if err != nil {
								resultCallback(err.Error(), result)
								return
							}

							resultCallback(nil, result)
						}

						if handler.Concurrent {
							go func() {
								defer log.RecoverAndLog()
								execHandler()
							}()
							return
						}

						execHandler()
					}

					go func() {
						defer log.RecoverAndLog()
						for data := range d.SendChan {
							log.Debug("Write", client.RoutingKey, data)
							err := publishChannel.Publish("broker", client.RoutingKey, false, false, amqp.Publishing{Body: data})
							if err != nil {
								log.LogError(err, 0)
							}
						}
					}()

					d.Send("ready", "kite-"+k.Name)

					for {
						select {
						case message, ok := <-channel:
							if !ok {
								return
							}
							log.Debug("Read", client.RoutingKey, message)
							d.ProcessMessage(message)
						case <-time.After(24 * time.Hour):
							timeoutChannel <- client.RoutingKey
						}
					}
				}()

			case "auth.leave":
				var client struct {
					RoutingKey string
				}
				err := json.Unmarshal(message.Body, &client)
				if err != nil || client.RoutingKey == "" {
					log.Err("Invalid auth.leave message.", message.Body)
					continue
				}

				channel, found := routeMap[client.RoutingKey]
				if found {
					close(channel)
					delete(routeMap, client.RoutingKey)
				}

			default:
				channel, found := routeMap[message.RoutingKey]
				if found {
					select {
					case channel <- message.Body:
						// successful
					default:
						close(channel)
						delete(routeMap, message.RoutingKey)
						log.Warn("Dropped client because of message buffer overflow.")
					}
				}
			}

		case routingKey := <-timeoutChannel:
			channel, found := routeMap[routingKey]
			if found {
				close(channel)
				delete(routeMap, routingKey)
				log.Warn("Dropped client because of fallback session timeout.")
			}

		case <-sigtermChannel:
			log.Info("Received TERM signal. Beginning shutdown...")
			lifecycle.BeginShutdown()
			consumeChannel.Close()
		}
	}
}

type Session struct {
	Username     string
	Alive        bool
	onDisconnect []func()
}

func NewSession(username string) *Session {
	return &Session{
		Username: username,
		Alive:    true,
	}
}

func (session *Session) OnDisconnect(f func()) {
	session.onDisconnect = append(session.onDisconnect, f)
}

func (session *Session) Close() {
	session.Alive = false
	for _, f := range session.onDisconnect {
		f()
	}
	session.onDisconnect = nil
}
<|endoftext|>


<commit_before>package typemap_test

// TODO(adonovan):
// - test use of explicit hasher across two maps.
// - test hashcodes are consistent with equals for a range of types
//   (e.g. all types generated by type-checking some body of real code).

import (
	"go/ast"
	"testing"

	"code.google.com/p/go.tools/go/types"
	"code.google.com/p/go.tools/go/types/typemap"
)

var (
	tStr      = types.Typ[types.String]       // string
	tPStr1    = types.NewPointer(tStr)        // *string
	tPStr2    = types.NewPointer(tStr)        // *string, again
	tInt      = types.Typ[types.Int]          // int
	tChanInt1 = types.NewChan(ast.RECV, tInt) // <-chan int
	tChanInt2 = types.NewChan(ast.RECV, tInt) // <-chan int, again
)

func checkEqualButNotIdentical(t *testing.T, x, y types.Type, comment string) {
	if !types.IsIdentical(x, y) {
		t.Errorf("%s: not equal: %s, %s", comment, x, y)
	}
	if x == y {
		t.Errorf("%s: identical: %p, %p", comment, x, y)
	}
}

func TestAxioms(t *testing.T) {
	checkEqualButNotIdentical(t, tPStr1, tPStr2, "tPstr{1,2}")
	checkEqualButNotIdentical(t, tChanInt1, tChanInt2, "tChanInt{1,2}")
}

func TestTypeMap(t *testing.T) {
	var tmap *typemap.M

	// All methods but Set are safe on on (*T)(nil).
	tmap.Len()
	tmap.At(tPStr1)
	tmap.Delete(tPStr1)
	tmap.KeysString()
	tmap.String()

	tmap = new(typemap.M)

	// Length of empty map.
	if l := tmap.Len(); l != 0 {
		t.Errorf("Len() on empty typemap: got %d, want 0", l)
	}
	// At of missing key.
	if v := tmap.At(tPStr1); v != nil {
		t.Errorf("At() on empty typemap: got %v, want nil", v)
	}
	// Deletion of missing key.
	if tmap.Delete(tPStr1) {
		t.Errorf("Delete() on empty typemap: got true, want false")
	}
	// Set of new key.
	if prev := tmap.Set(tPStr1, "*string"); prev != nil {
		t.Errorf("Set() on empty map returned non-nil previous value %s", prev)
	}

	// Now: {*string: "*string"}

	// Length of non-empty map.
	if l := tmap.Len(); l != 1 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// At via insertion key.
	if v := tmap.At(tPStr1); v != "*string" {
		t.Errorf("At(): got %q, want \"*string\"", v)
	}
	// At via equal key.
	if v := tmap.At(tPStr2); v != "*string" {
		t.Errorf("At(): got %q, want \"*string\"", v)
	}
	// Iteration over sole entry.
	tmap.Iterate(func(key types.Type, value interface{}) {
		if key != tPStr1 {
			t.Errorf("Iterate: key: got %s, want %s", key, tPStr1)
		}
		if want := "*string"; value != want {
			t.Errorf("Iterate: value: got %s, want %s", value, want)
		}
	})

	// Setion with key equal to present one.
	if prev := tmap.Set(tPStr2, "*string again"); prev != "*string" {
		t.Errorf("Set() previous value: got %s, want \"*string\"", prev)
	}

	// Setion of another association.
	if prev := tmap.Set(tChanInt1, "<-chan int"); prev != nil {
		t.Errorf("Set() previous value: got %s, want nil", prev)
	}

	// Now: {*string: "*string again", <-chan int: "<-chan int"}

	want1 := "{*string: \"*string again\", <-chan int: \"<-chan int\"}"
	want2 := "{<-chan int: \"<-chan int\", *string: \"*string again\"}"
	if s := tmap.String(); s != want1 && s != want2 {
		t.Errorf("String(): got %s, want %s", s, want1)
	}

	want1 = "{*string, <-chan int}"
	want2 = "{<-chan int, *string}"
	if s := tmap.KeysString(); s != want1 && s != want2 {
		t.Errorf("KeysString(): got %s, want %s", s, want1)
	}

	// Keys().
	I := types.IsIdentical
	switch k := tmap.Keys(); {
	case I(k[0], tChanInt1) && I(k[1], tPStr1): // ok
	case I(k[1], tChanInt1) && I(k[0], tPStr1): // ok
	default:
		t.Errorf("Keys(): got %v, want %s", k, want2)
	}

	if l := tmap.Len(); l != 2 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// At via original key.
	if v := tmap.At(tPStr1); v != "*string again" {
		t.Errorf("At(): got %q, want \"*string again\"", v)
	}
	hamming := 1
	tmap.Iterate(func(key types.Type, value interface{}) {
		switch {
		case I(key, tChanInt1):
			hamming *= 2 // ok
		case I(key, tPStr1):
			hamming *= 3 // ok
		}
	})
	if hamming != 6 {
		t.Errorf("Iterate: hamming: got %d, want %d", hamming, 6)
	}

	if v := tmap.At(tChanInt2); v != "<-chan int" {
		t.Errorf("At(): got %q, want \"<-chan int\"", v)
	}
	// Deletion with key equal to present one.
	if !tmap.Delete(tChanInt2) {
		t.Errorf("Delete() of existing key: got false, want true")
	}

	// Now: {*string: "*string again"}

	if l := tmap.Len(); l != 1 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// Deletion again.
	if !tmap.Delete(tPStr2) {
		t.Errorf("Delete() of existing key: got false, want true")
	}

	// Now: {}

	if l := tmap.Len(); l != 0 {
		t.Errorf("Len(): got %d, want %d", l, 0)
	}
	if s := tmap.String(); s != "{}" {
		t.Errorf("Len(): got %q, want %q", s, "")
	}
}
<commit_msg>go.tools/go/types/typemap: fix formatting verb<commit_after>package typemap_test

// TODO(adonovan):
// - test use of explicit hasher across two maps.
// - test hashcodes are consistent with equals for a range of types
//   (e.g. all types generated by type-checking some body of real code).

import (
	"go/ast"
	"testing"

	"code.google.com/p/go.tools/go/types"
	"code.google.com/p/go.tools/go/types/typemap"
)

var (
	tStr      = types.Typ[types.String]       // string
	tPStr1    = types.NewPointer(tStr)        // *string
	tPStr2    = types.NewPointer(tStr)        // *string, again
	tInt      = types.Typ[types.Int]          // int
	tChanInt1 = types.NewChan(ast.RECV, tInt) // <-chan int
	tChanInt2 = types.NewChan(ast.RECV, tInt) // <-chan int, again
)

func checkEqualButNotIdentical(t *testing.T, x, y types.Type, comment string) {
	if !types.IsIdentical(x, y) {
		t.Errorf("%s: not equal: %s, %s", comment, x, y)
	}
	if x == y {
		t.Errorf("%s: identical: %v, %v", comment, x, y)
	}
}

func TestAxioms(t *testing.T) {
	checkEqualButNotIdentical(t, tPStr1, tPStr2, "tPstr{1,2}")
	checkEqualButNotIdentical(t, tChanInt1, tChanInt2, "tChanInt{1,2}")
}

func TestTypeMap(t *testing.T) {
	var tmap *typemap.M

	// All methods but Set are safe on on (*T)(nil).
	tmap.Len()
	tmap.At(tPStr1)
	tmap.Delete(tPStr1)
	tmap.KeysString()
	tmap.String()

	tmap = new(typemap.M)

	// Length of empty map.
	if l := tmap.Len(); l != 0 {
		t.Errorf("Len() on empty typemap: got %d, want 0", l)
	}
	// At of missing key.
	if v := tmap.At(tPStr1); v != nil {
		t.Errorf("At() on empty typemap: got %v, want nil", v)
	}
	// Deletion of missing key.
	if tmap.Delete(tPStr1) {
		t.Errorf("Delete() on empty typemap: got true, want false")
	}
	// Set of new key.
	if prev := tmap.Set(tPStr1, "*string"); prev != nil {
		t.Errorf("Set() on empty map returned non-nil previous value %s", prev)
	}

	// Now: {*string: "*string"}

	// Length of non-empty map.
	if l := tmap.Len(); l != 1 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// At via insertion key.
	if v := tmap.At(tPStr1); v != "*string" {
		t.Errorf("At(): got %q, want \"*string\"", v)
	}
	// At via equal key.
	if v := tmap.At(tPStr2); v != "*string" {
		t.Errorf("At(): got %q, want \"*string\"", v)
	}
	// Iteration over sole entry.
	tmap.Iterate(func(key types.Type, value interface{}) {
		if key != tPStr1 {
			t.Errorf("Iterate: key: got %s, want %s", key, tPStr1)
		}
		if want := "*string"; value != want {
			t.Errorf("Iterate: value: got %s, want %s", value, want)
		}
	})

	// Setion with key equal to present one.
	if prev := tmap.Set(tPStr2, "*string again"); prev != "*string" {
		t.Errorf("Set() previous value: got %s, want \"*string\"", prev)
	}

	// Setion of another association.
	if prev := tmap.Set(tChanInt1, "<-chan int"); prev != nil {
		t.Errorf("Set() previous value: got %s, want nil", prev)
	}

	// Now: {*string: "*string again", <-chan int: "<-chan int"}

	want1 := "{*string: \"*string again\", <-chan int: \"<-chan int\"}"
	want2 := "{<-chan int: \"<-chan int\", *string: \"*string again\"}"
	if s := tmap.String(); s != want1 && s != want2 {
		t.Errorf("String(): got %s, want %s", s, want1)
	}

	want1 = "{*string, <-chan int}"
	want2 = "{<-chan int, *string}"
	if s := tmap.KeysString(); s != want1 && s != want2 {
		t.Errorf("KeysString(): got %s, want %s", s, want1)
	}

	// Keys().
	I := types.IsIdentical
	switch k := tmap.Keys(); {
	case I(k[0], tChanInt1) && I(k[1], tPStr1): // ok
	case I(k[1], tChanInt1) && I(k[0], tPStr1): // ok
	default:
		t.Errorf("Keys(): got %v, want %s", k, want2)
	}

	if l := tmap.Len(); l != 2 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// At via original key.
	if v := tmap.At(tPStr1); v != "*string again" {
		t.Errorf("At(): got %q, want \"*string again\"", v)
	}
	hamming := 1
	tmap.Iterate(func(key types.Type, value interface{}) {
		switch {
		case I(key, tChanInt1):
			hamming *= 2 // ok
		case I(key, tPStr1):
			hamming *= 3 // ok
		}
	})
	if hamming != 6 {
		t.Errorf("Iterate: hamming: got %d, want %d", hamming, 6)
	}

	if v := tmap.At(tChanInt2); v != "<-chan int" {
		t.Errorf("At(): got %q, want \"<-chan int\"", v)
	}
	// Deletion with key equal to present one.
	if !tmap.Delete(tChanInt2) {
		t.Errorf("Delete() of existing key: got false, want true")
	}

	// Now: {*string: "*string again"}

	if l := tmap.Len(); l != 1 {
		t.Errorf("Len(): got %d, want 1", l)
	}
	// Deletion again.
	if !tmap.Delete(tPStr2) {
		t.Errorf("Delete() of existing key: got false, want true")
	}

	// Now: {}

	if l := tmap.Len(); l != 0 {
		t.Errorf("Len(): got %d, want %d", l, 0)
	}
	if s := tmap.String(); s != "{}" {
		t.Errorf("Len(): got %q, want %q", s, "")
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package archiver

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/luci/luci-go/client/internal/common"
	"github.com/luci/luci-go/client/isolatedclient"
	"github.com/luci/luci-go/common/isolated"
)

// Future is the future of a file pushed to be hashed.
type Future interface {
	// WaitForHashed hangs until the item hash is known.
	WaitForHashed()
	// Error returns any error that occured for this item if any.
	Error() error
	// Digest returns the calculated digest once calculated, empty otherwise.
	Digest() isolated.HexDigest
}

// Archiver is an high level interface to an isolatedclient.IsolateServer.
type Archiver interface {
	io.Closer
	common.Cancelable
	Push(displayName string, src io.ReadSeeker) Future
	PushFile(path string) Future
	Stats() *Stats
}

// UploadStat is the statistic for a single upload.
type UploadStat struct {
	Duration time.Duration
	Size     int64
}

// Statistics from the Archiver.
type Stats struct {
	Hits   []int64
	Misses []int64
	Pushed []*UploadStat
}

func (s *Stats) TotalHits() int64 {
	out := int64(0)
	for _, i := range s.Hits {
		out += i
	}
	return out
}

func (s *Stats) TotalMisses() int64 {
	out := int64(0)
	for _, i := range s.Misses {
		out += i
	}
	return out
}

func (s *Stats) TotalPushed() int64 {
	out := int64(0)
	for _, i := range s.Pushed {
		out += i.Size
	}
	return out
}

func (s *Stats) deepCopy() *Stats {
	out := &Stats{}
	out.Hits = make([]int64, len(s.Hits))
	copy(out.Hits, s.Hits)
	out.Misses = make([]int64, len(s.Misses))
	copy(out.Misses, s.Misses)
	out.Pushed = make([]*UploadStat, len(s.Pushed))
	copy(out.Pushed, s.Pushed)
	return out
}

// New returns a thread-safe Archiver instance.
//
// TODO(maruel): Cache hashes and server cache presence.
func New(is isolatedclient.IsolateServer) Archiver {
	a := &archiver{
		canceler:              common.NewCanceler(),
		is:                    is,
		maxConcurrentHash:     5,
		maxConcurrentContains: 16,
		maxConcurrentUpload:   8,
		containsBatchingDelay: 100 * time.Millisecond,
		filesToHash:           make(chan *archiverItem, 10240),
		itemsToLookup:         make(chan *archiverItem, 10240),
		itemsToUpload:         make(chan *archiverItem, 10240),
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.hashLoop()
	}()

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.containsLoop()
	}()

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.uploadLoop()
	}()
	return a
}

// Private details.

// archiverItem is an item to process. Implements Future.
//
// It is caried over from pipeline stage to stage to do processing on it.
type archiverItem struct {
	// Immutable.
	path        string         // Set when source is a file on disk
	displayName string         // Name to use to qualify this item
	wgHashed    sync.WaitGroup // Released once .digestItem.Digest is set

	// Mutable.
	lock       sync.Mutex
	err        error               // Item specific error
	digestItem isolated.DigestItem // Mutated by hashLoop(), used by doContains()

	// Mutable but not accessible externally.
	src   io.ReadSeeker             // Source of data
	state *isolatedclient.PushState // Server-side push state for cache miss
}

func newArchiverItem(path, displayName string, src io.ReadSeeker) *archiverItem {
	i := &archiverItem{path: path, displayName: displayName, src: src}
	i.wgHashed.Add(1)
	return i
}

func (i *archiverItem) WaitForHashed() {
	i.wgHashed.Wait()
}

func (i *archiverItem) Error() error {
	i.lock.Lock()
	defer i.lock.Unlock()
	return i.err
}

func (i *archiverItem) Digest() isolated.HexDigest {
	i.lock.Lock()
	defer i.lock.Unlock()
	return i.digestItem.Digest
}

func (i *archiverItem) setErr(err error) {
	if err == nil {
		panic("internal error")
	}
	i.lock.Lock()
	defer i.lock.Unlock()
	if i.err == nil {
		i.err = err
	}
	// TODO(maruel): Support Close().
	i.src = nil
}

func (i *archiverItem) calcDigest() error {
	defer i.wgHashed.Done()
	var d isolated.DigestItem
	if i.path != "" {
		// Open and hash the file.
		var err error
		if d, err = isolated.HashFile(i.path); err != nil {
			i.setErr(err)
			return fmt.Errorf("hash(%s) failed: %s\n", i.displayName, err)
		}
	} else {
		// Use src instead.
		h := isolated.GetHash()
		size, err := io.Copy(h, i.src)
		if err != nil {
			i.setErr(err)
			return fmt.Errorf("read(%s) failed: %s\n", i.displayName, err)
		}
		_, err = i.src.Seek(0, os.SEEK_SET)
		if err != nil {
			i.setErr(err)
			return fmt.Errorf("seek(%s) failed: %s\n", i.displayName, err)
		}
		digest := isolated.HexDigest(hex.EncodeToString(h.Sum(nil)))
		d = isolated.DigestItem{digest, true, size}
	}
	i.lock.Lock()
	defer i.lock.Unlock()
	i.digestItem = d
	return nil
}

// archiver archives content to an Isolate server.
//
// Uses a 3 stages pipeline:
// - Hashing files concurrently.
// - Batched cache hit lookups on the server.
// - Uploading cache misses concurrently.
//
// TODO(maruel): Add another pipeline stage to do a local hash cache hit to
// skip hashing the files at all.
type archiver struct {
	// Immutable.
	is                    isolatedclient.IsolateServer
	maxConcurrentHash     int                // Stage 1
	maxConcurrentContains int                // Stage 2
	maxConcurrentUpload   int                // Stage 3
	containsBatchingDelay time.Duration      // Used by stage 2
	filesToHash           chan *archiverItem // Stage 1
	itemsToLookup         chan *archiverItem // Stage 2
	itemsToUpload         chan *archiverItem // Stage 3
	wg                    sync.WaitGroup
	canceler              common.Canceler

	// Mutable.
	lock  sync.Mutex
	stats Stats
}

// Close waits for all pending files to be done.
func (a *archiver) Close() error {
	// This is done so asynchronously calling push() won't crash.
	a.lock.Lock()
	close(a.filesToHash)
	a.filesToHash = nil
	a.lock.Unlock()

	a.wg.Wait()
	_ = a.canceler.Close()
	return a.CancelationReason()
}

func (a *archiver) Cancel(reason error) {
	a.canceler.Cancel(reason)
}

func (a *archiver) CancelationReason() error {
	return a.canceler.CancelationReason()
}

func (a *archiver) Push(displayName string, src io.ReadSeeker) Future {
	i := newArchiverItem("", displayName, src)
	if pos, err := i.src.Seek(0, os.SEEK_SET); pos != 0 || err != nil {
		i.err = errors.New("must use buffer set at offset 0")
		return i
	}
	return a.push(i)
}

func (a *archiver) PushFile(path string) Future {
	return a.push(newArchiverItem(path, path, nil))
}

func (a *archiver) Stats() *Stats {
	a.lock.Lock()
	defer a.lock.Unlock()
	return a.stats.deepCopy()
}

func (a *archiver) push(item *archiverItem) Future {
	a.lock.Lock()
	defer a.lock.Unlock()
	if a.filesToHash == nil {
		return nil
	}
	a.filesToHash <- item
	return item
}

// hashLoop is stage 1.
func (a *archiver) hashLoop() {
	defer close(a.itemsToLookup)

	pool := common.NewGoroutinePool(a.maxConcurrentHash, a.canceler)

	a.lock.Lock()
	filesToHash := a.filesToHash
	a.lock.Unlock()
	if filesToHash == nil {
		// This means Close() was called really too quickly.
		return
	}
	for file := range filesToHash {
		item := file
		pool.Schedule(func() {
			if err := item.calcDigest(); err != nil {
				a.Cancel(err)
				return
			}
			a.itemsToLookup <- item
		})
	}

	_ = pool.Wait()
}

// containsLoop is stage 2.
func (a *archiver) containsLoop() {
	defer close(a.itemsToUpload)

	pool := common.NewGoroutinePool(a.maxConcurrentContains, a.canceler)

	items := []*archiverItem{}
	never := make(<-chan time.Time)
	timer := never
	loop := true
	for loop {
		select {
		case <-timer:
			batch := items
			pool.Schedule(func() {
				a.doContains(batch)
			})
			items = []*archiverItem{}
			timer = never

		case item, ok := <-a.itemsToLookup:
			if !ok {
				loop = false
				break
			}
			items = append(items, item)
			if timer == never {
				timer = time.After(a.containsBatchingDelay)
			}
		}
	}

	if len(items) != 0 {
		batch := items
		pool.Schedule(func() {
			a.doContains(batch)
		})
	}
	_ = pool.Wait()
}

// uploadLoop is stage 3.
func (a *archiver) uploadLoop() {
	pool := common.NewGoroutinePool(a.maxConcurrentUpload, a.canceler)

	for state := range a.itemsToUpload {
		item := state
		pool.Schedule(func() {
			a.doUpload(item)
		})
	}
	_ = pool.Wait()
}

// doContains is scaled by stage 2.
func (a *archiver) doContains(items []*archiverItem) {
	tmp := make([]*isolated.DigestItem, len(items))
	// No need to lock each item at that point, no mutation occurs on
	// archiverItem.digestItem after stage 1.
	for i, item := range items {
		tmp[i] = &item.digestItem
	}
	states, err := a.is.Contains(tmp)
	if err != nil {
		err = fmt.Errorf("contains(%d) failed: %s", len(items), err)
		a.Cancel(err)
		for _, item := range items {
			item.setErr(err)
		}
		return
	}
	for index, state := range states {
		size := items[index].digestItem.Size
		if state == nil {
			a.lock.Lock()
			a.stats.Hits = append(a.stats.Hits, size)
			a.lock.Unlock()
		} else {
			a.lock.Lock()
			a.stats.Misses = append(a.stats.Misses, size)
			a.lock.Unlock()
			items[index].state = state
			a.itemsToUpload <- items[index]
		}
	}
	log.Printf("Looked up %d items\n", len(items))
}

// doUpload is scaled by stage 3.
func (a *archiver) doUpload(item *archiverItem) {
	var src io.Reader
	if item.src == nil {
		f, err := os.Open(item.path)
		if err != nil {
			a.Cancel(err)
			item.setErr(err)
			return
		}
		defer f.Close()
		src = f
	} else {
		src = item.src
		item.src = nil
	}
	start := time.Now()
	if err := a.is.Push(item.state, src); err != nil {
		err = fmt.Errorf("push(%s) failed: %s\n", item.path, err)
		a.Cancel(err)
		item.setErr(err)
	}
	duration := time.Now().Sub(start)
	u := &UploadStat{duration, item.digestItem.Size}
	a.lock.Lock()
	a.stats.Pushed = append(a.stats.Pushed, u)
	a.lock.Unlock()
	log.Printf("Uploaded %.1fkb\n", float64(item.digestItem.Size)/1024.)
}
<commit_msg>Increase concurrent cache lookup from 16 to 64 and limit each batch to 50 items.<commit_after>// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package archiver

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	"github.com/luci/luci-go/client/internal/common"
	"github.com/luci/luci-go/client/isolatedclient"
	"github.com/luci/luci-go/common/isolated"
)

// Future is the future of a file pushed to be hashed.
type Future interface {
	// WaitForHashed hangs until the item hash is known.
	WaitForHashed()
	// Error returns any error that occured for this item if any.
	Error() error
	// Digest returns the calculated digest once calculated, empty otherwise.
	Digest() isolated.HexDigest
}

// Archiver is an high level interface to an isolatedclient.IsolateServer.
type Archiver interface {
	io.Closer
	common.Cancelable
	Push(displayName string, src io.ReadSeeker) Future
	PushFile(path string) Future
	Stats() *Stats
}

// UploadStat is the statistic for a single upload.
type UploadStat struct {
	Duration time.Duration
	Size     int64
}

// Statistics from the Archiver.
type Stats struct {
	Hits   []int64
	Misses []int64
	Pushed []*UploadStat
}

func (s *Stats) TotalHits() int64 {
	out := int64(0)
	for _, i := range s.Hits {
		out += i
	}
	return out
}

func (s *Stats) TotalMisses() int64 {
	out := int64(0)
	for _, i := range s.Misses {
		out += i
	}
	return out
}

func (s *Stats) TotalPushed() int64 {
	out := int64(0)
	for _, i := range s.Pushed {
		out += i.Size
	}
	return out
}

func (s *Stats) deepCopy() *Stats {
	out := &Stats{}
	out.Hits = make([]int64, len(s.Hits))
	copy(out.Hits, s.Hits)
	out.Misses = make([]int64, len(s.Misses))
	copy(out.Misses, s.Misses)
	out.Pushed = make([]*UploadStat, len(s.Pushed))
	copy(out.Pushed, s.Pushed)
	return out
}

// New returns a thread-safe Archiver instance.
//
// TODO(maruel): Cache hashes and server cache presence.
func New(is isolatedclient.IsolateServer) Archiver {
	a := &archiver{
		canceler:              common.NewCanceler(),
		is:                    is,
		maxConcurrentHash:     5,
		maxConcurrentContains: 64,
		maxConcurrentUpload:   8,
		containsBatchingDelay: 100 * time.Millisecond,
		containsBatchSize:     50,
		filesToHash:           make(chan *archiverItem, 10240),
		itemsToLookup:         make(chan *archiverItem, 10240),
		itemsToUpload:         make(chan *archiverItem, 10240),
	}

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.hashLoop()
	}()

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.containsLoop()
	}()

	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		a.uploadLoop()
	}()
	return a
}

// Private details.

// archiverItem is an item to process. Implements Future.
//
// It is caried over from pipeline stage to stage to do processing on it.
type archiverItem struct {
	// Immutable.
	path        string         // Set when source is a file on disk
	displayName string         // Name to use to qualify this item
	wgHashed    sync.WaitGroup // Released once .digestItem.Digest is set

	// Mutable.
	lock       sync.Mutex
	err        error               // Item specific error
	digestItem isolated.DigestItem // Mutated by hashLoop(), used by doContains()

	// Mutable but not accessible externally.
	src   io.ReadSeeker             // Source of data
	state *isolatedclient.PushState // Server-side push state for cache miss
}

func newArchiverItem(path, displayName string, src io.ReadSeeker) *archiverItem {
	i := &archiverItem{path: path, displayName: displayName, src: src}
	i.wgHashed.Add(1)
	return i
}

func (i *archiverItem) WaitForHashed() {
	i.wgHashed.Wait()
}

func (i *archiverItem) Error() error {
	i.lock.Lock()
	defer i.lock.Unlock()
	return i.err
}

func (i *archiverItem) Digest() isolated.HexDigest {
	i.lock.Lock()
	defer i.lock.Unlock()
	return i.digestItem.Digest
}

func (i *archiverItem) setErr(err error) {
	if err == nil {
		panic("internal error")
	}
	i.lock.Lock()
	defer i.lock.Unlock()
	if i.err == nil {
		i.err = err
	}
	// TODO(maruel): Support Close().
	i.src = nil
}

func (i *archiverItem) calcDigest() error {
	defer i.wgHashed.Done()
	var d isolated.DigestItem
	if i.path != "" {
		// Open and hash the file.
		var err error
		if d, err = isolated.HashFile(i.path); err != nil {
			i.setErr(err)
			return fmt.Errorf("hash(%s) failed: %s\n", i.displayName, err)
		}
	} else {
		// Use src instead.
		h := isolated.GetHash()
		size, err := io.Copy(h, i.src)
		if err != nil {
			i.setErr(err)
			return fmt.Errorf("read(%s) failed: %s\n", i.displayName, err)
		}
		_, err = i.src.Seek(0, os.SEEK_SET)
		if err != nil {
			i.setErr(err)
			return fmt.Errorf("seek(%s) failed: %s\n", i.displayName, err)
		}
		digest := isolated.HexDigest(hex.EncodeToString(h.Sum(nil)))
		d = isolated.DigestItem{digest, true, size}
	}
	i.lock.Lock()
	defer i.lock.Unlock()
	i.digestItem = d
	return nil
}

// archiver archives content to an Isolate server.
//
// Uses a 3 stages pipeline:
// - Hashing files concurrently.
// - Batched cache hit lookups on the server.
// - Uploading cache misses concurrently.
//
// TODO(maruel): Add another pipeline stage to do a local hash cache hit to
// skip hashing the files at all.
type archiver struct {
	// Immutable.
	is                    isolatedclient.IsolateServer
	maxConcurrentHash     int                // Stage 1
	maxConcurrentContains int                // Stage 2
	maxConcurrentUpload   int                // Stage 3
	containsBatchingDelay time.Duration      // Used by stage 2
	containsBatchSize     int                // Used by stage 2
	filesToHash           chan *archiverItem // Stage 1
	itemsToLookup         chan *archiverItem // Stage 2
	itemsToUpload         chan *archiverItem // Stage 3
	wg                    sync.WaitGroup
	canceler              common.Canceler

	// Mutable.
	lock  sync.Mutex
	stats Stats
}

// Close waits for all pending files to be done.
func (a *archiver) Close() error {
	// This is done so asynchronously calling push() won't crash.
	a.lock.Lock()
	close(a.filesToHash)
	a.filesToHash = nil
	a.lock.Unlock()

	a.wg.Wait()
	_ = a.canceler.Close()
	return a.CancelationReason()
}

func (a *archiver) Cancel(reason error) {
	a.canceler.Cancel(reason)
}

func (a *archiver) CancelationReason() error {
	return a.canceler.CancelationReason()
}

func (a *archiver) Push(displayName string, src io.ReadSeeker) Future {
	i := newArchiverItem("", displayName, src)
	if pos, err := i.src.Seek(0, os.SEEK_SET); pos != 0 || err != nil {
		i.err = errors.New("must use buffer set at offset 0")
		return i
	}
	return a.push(i)
}

func (a *archiver) PushFile(path string) Future {
	return a.push(newArchiverItem(path, path, nil))
}

func (a *archiver) Stats() *Stats {
	a.lock.Lock()
	defer a.lock.Unlock()
	return a.stats.deepCopy()
}

func (a *archiver) push(item *archiverItem) Future {
	a.lock.Lock()
	defer a.lock.Unlock()
	if a.filesToHash == nil {
		return nil
	}
	a.filesToHash <- item
	return item
}

// hashLoop is stage 1.
func (a *archiver) hashLoop() {
	defer close(a.itemsToLookup)

	pool := common.NewGoroutinePool(a.maxConcurrentHash, a.canceler)

	a.lock.Lock()
	filesToHash := a.filesToHash
	a.lock.Unlock()
	if filesToHash == nil {
		// This means Close() was called really too quickly.
		return
	}
	for file := range filesToHash {
		item := file
		pool.Schedule(func() {
			if err := item.calcDigest(); err != nil {
				a.Cancel(err)
				return
			}
			a.itemsToLookup <- item
		})
	}

	_ = pool.Wait()
}

// containsLoop is stage 2.
func (a *archiver) containsLoop() {
	defer close(a.itemsToUpload)

	pool := common.NewGoroutinePool(a.maxConcurrentContains, a.canceler)

	items := []*archiverItem{}
	never := make(<-chan time.Time)
	timer := never
	loop := true
	for loop {
		select {
		case <-timer:
			batch := items
			pool.Schedule(func() {
				a.doContains(batch)
			})
			items = []*archiverItem{}
			timer = never

		case item, ok := <-a.itemsToLookup:
			if !ok {
				loop = false
				break
			}
			items = append(items, item)
			if len(items) == a.containsBatchSize {
				batch := items
				pool.Schedule(func() {
					a.doContains(batch)
				})
				items = []*archiverItem{}
				timer = never
			} else if timer == never {
				timer = time.After(a.containsBatchingDelay)
			}
		}
	}

	if len(items) != 0 {
		batch := items
		pool.Schedule(func() {
			a.doContains(batch)
		})
	}
	_ = pool.Wait()
}

// uploadLoop is stage 3.
func (a *archiver) uploadLoop() {
	pool := common.NewGoroutinePool(a.maxConcurrentUpload, a.canceler)

	for state := range a.itemsToUpload {
		item := state
		pool.Schedule(func() {
			a.doUpload(item)
		})
	}
	_ = pool.Wait()
}

// doContains is scaled by stage 2.
func (a *archiver) doContains(items []*archiverItem) {
	tmp := make([]*isolated.DigestItem, len(items))
	// No need to lock each item at that point, no mutation occurs on
	// archiverItem.digestItem after stage 1.
	for i, item := range items {
		tmp[i] = &item.digestItem
	}
	states, err := a.is.Contains(tmp)
	if err != nil {
		err = fmt.Errorf("contains(%d) failed: %s", len(items), err)
		a.Cancel(err)
		for _, item := range items {
			item.setErr(err)
		}
		return
	}
	for index, state := range states {
		size := items[index].digestItem.Size
		if state == nil {
			a.lock.Lock()
			a.stats.Hits = append(a.stats.Hits, size)
			a.lock.Unlock()
		} else {
			a.lock.Lock()
			a.stats.Misses = append(a.stats.Misses, size)
			a.lock.Unlock()
			items[index].state = state
			a.itemsToUpload <- items[index]
		}
	}
	log.Printf("Looked up %d items\n", len(items))
}

// doUpload is scaled by stage 3.
func (a *archiver) doUpload(item *archiverItem) {
	var src io.Reader
	if item.src == nil {
		f, err := os.Open(item.path)
		if err != nil {
			a.Cancel(err)
			item.setErr(err)
			return
		}
		defer f.Close()
		src = f
	} else {
		src = item.src
		item.src = nil
	}
	start := time.Now()
	if err := a.is.Push(item.state, src); err != nil {
		err = fmt.Errorf("push(%s) failed: %s\n", item.path, err)
		a.Cancel(err)
		item.setErr(err)
	}
	duration := time.Now().Sub(start)
	u := &UploadStat{duration, item.digestItem.Size}
	a.lock.Lock()
	a.stats.Pushed = append(a.stats.Pushed, u)
	a.lock.Unlock()
	log.Printf("Uploaded %.1fkb\n", float64(item.digestItem.Size)/1024.)
}
<|endoftext|>