<commit_before>package oauth

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"
	"net/http/httputil"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
	oauthv1 "github.com/openshift/api/oauth/v1"
	userv1 "github.com/openshift/api/user/v1"
	exutil "github.com/openshift/origin/test/extended/util"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/kubernetes/test/e2e/framework"
)

var _ = g.Describe("[sig-auth][Feature:OAuthServer]", func() {
	defer g.GinkgoRecover()
	oc := exutil.NewCLI("oauthclient-secret-with-plus")
	ctx := context.Background()

	g.Describe("ClientSecretWithPlus", func() {
		g.It(fmt.Sprintf("should create oauthclient"), func() {
			g.By("create oauth client")
			oauthClient, err := oc.AdminOauthClient().OauthV1().OAuthClients().Create(ctx, &oauthv1.OAuthClient{
				ObjectMeta: metav1.ObjectMeta{
					Name: "oauth-client-with-plus",
				},
				Secret:            "secret+with+plus",
				RedirectURIs:      []string{"https://www.google.com"},
				GrantMethod:       oauthv1.GrantHandlerAuto,
				ScopeRestrictions: []oauthv1.ScopeRestriction{{ExactValues: []string{"user:full"}}},
			}, metav1.CreateOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())
			oc.AddResourceToDelete(oauthv1.GroupVersion.WithResource("oauthclients"), oauthClient)

			g.By("create synthetic identity, user, and binding")

			user, err := oc.AdminUserClient().UserV1().Users().Create(ctx, &userv1.User{
				TypeMeta: metav1.TypeMeta{},
				ObjectMeta: metav1.ObjectMeta{
					Name: "user-for-plus",
				},
				FullName: "fake user",
			}, metav1.CreateOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())
			oc.AddResourceToDelete(userv1.GroupVersion.WithResource("users"), user)

			oauthAuthorizeToken, err := oc.AdminOauthClient().OauthV1().OAuthAuthorizeTokens().Create(ctx, &oauthv1.OAuthAuthorizeToken{
				ObjectMeta: metav1.ObjectMeta{
					Name: "oauth-client-with-plus-with-more-than-thirty-two-characters-in-this-very-long-name",
				},
				ClientName:          oauthClient.Name,
				ExpiresIn:           100000000,
				RedirectURI:         "https://www.google.com",
				Scopes:              []string{"user:full"},
				UserName:            user.Name,
				UserUID:             string(user.UID),
				CodeChallenge:       "",
				CodeChallengeMethod: "",
			}, metav1.CreateOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())
			oc.AddResourceToDelete(oauthv1.GroupVersion.WithResource("oauthauthorizetokens"), oauthAuthorizeToken)

			g.By("querying for a token")
			oauthRoute, err := oc.AdminRouteClient().RouteV1().Routes("openshift-authentication").Get(ctx, "oauth-openshift", metav1.GetOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())

			tokenPost, err := http.NewRequest("POST", "https://"+oauthRoute.Status.Ingress[0].Host+"/oauth/token?"+
				"grant_type=authorization_code&"+
				"code=oauth-client-with-plus-with-more-than-thirty-two-characters-in-this-very-long-name&"+
				"client_id=oauth-client-with-plus&"+
				"client_secret=secret%2Bwith%2Bplus",
				nil)
			o.Expect(err).NotTo(o.HaveOccurred())
			requestDump, err := httputil.DumpRequest(tokenPost, true)
			o.Expect(err).NotTo(o.HaveOccurred())
			framework.Logf("%v", string(requestDump))

			// we don't really care if this URL is safe
			tr := &http.Transport{
				Proxy:           http.ProxyFromEnvironment,
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			}
			client := &http.Client{Transport: tr}

			tokenResponse, err := client.Do(tokenPost)
			o.Expect(err).NotTo(o.HaveOccurred())
			response, err := httputil.DumpResponse(tokenResponse, true)
			o.Expect(err).NotTo(o.HaveOccurred())
			framework.Logf("%v", string(response))
			o.Expect(tokenResponse.StatusCode).To(o.Equal(http.StatusOK))
		})
	})
})
<commit_msg>e2e: Add other client-authn methods to token request test<commit_after>package oauth

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"net/http"
	"net/http/httputil"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/kubernetes/test/e2e/framework"

	oauthv1 "github.com/openshift/api/oauth/v1"
	userv1 "github.com/openshift/api/user/v1"
	exutil "github.com/openshift/origin/test/extended/util"
)

const authzTokenName string = "oauth-client-with-plus-with-more-than-thirty-two-characters-in-this-very-long-name"

var _ = g.Describe("[sig-auth][Feature:OAuthServer]", func() {
	defer g.GinkgoRecover()
	oc := exutil.NewCLI("oauthclient-secret-with-plus")
	ctx := context.Background()

	g.Describe("ClientSecretWithPlus", func() {
		g.It(fmt.Sprintf("should create oauthclient"), func() {
			g.By("create oauth client")
			oauthClient, err := oc.AdminOauthClient().OauthV1().OAuthClients().Create(ctx, &oauthv1.OAuthClient{
				ObjectMeta: metav1.ObjectMeta{
					Name: "oauth-client-with-plus",
				},
				Secret:            "secret+with+plus",
				RedirectURIs:      []string{"https://www.google.com"},
				GrantMethod:       oauthv1.GrantHandlerAuto,
				ScopeRestrictions: []oauthv1.ScopeRestriction{{ExactValues: []string{"user:full"}}},
			}, metav1.CreateOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())
			oc.AddResourceToDelete(oauthv1.GroupVersion.WithResource("oauthclients"), oauthClient)

			g.By("create synthetic user")

			user, err := oc.AdminUserClient().UserV1().Users().Create(ctx, &userv1.User{
				TypeMeta: metav1.TypeMeta{},
				ObjectMeta: metav1.ObjectMeta{
					Name: "user-for-plus",
				},
				FullName: "fake user",
			}, metav1.CreateOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())
			oc.AddResourceToDelete(userv1.GroupVersion.WithResource("users"), user)

			oauthRoute, err := oc.AdminRouteClient().RouteV1().Routes("openshift-authentication").Get(ctx, "oauth-openshift", metav1.GetOptions{})
			o.Expect(err).NotTo(o.HaveOccurred())

			for _, createTokenReq := range []func(string) *http.Request{
				queryClientAuthRequest,
				bodyClientAuthRequest,
				headerClientAuthRequest,
			} {
				g.By("create synthetic authz token")
				oauthAuthorizeToken, err := oc.AdminOauthClient().OauthV1().OAuthAuthorizeTokens().Create(ctx, &oauthv1.OAuthAuthorizeToken{
					ObjectMeta: metav1.ObjectMeta{
						Name: authzTokenName,
					},
					ClientName:          oauthClient.Name,
					ExpiresIn:           100000000,
					RedirectURI:         "https://www.google.com",
					Scopes:              []string{"user:full"},
					UserName:            user.Name,
					UserUID:             string(user.UID),
					CodeChallenge:       "",
					CodeChallengeMethod: "",
				}, metav1.CreateOptions{})
				o.Expect(err).NotTo(o.HaveOccurred())
				oc.AddResourceToDelete(oauthv1.GroupVersion.WithResource("oauthauthorizetokens"), oauthAuthorizeToken)

				g.By("querying for a token")

				tokenReq := createTokenReq(oauthRoute.Status.Ingress[0].Host)
				tokenReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
				requestDump, err := httputil.DumpRequest(tokenReq, true)
				o.Expect(err).NotTo(o.HaveOccurred())
				framework.Logf("%v", string(requestDump))

				// we don't really care if this URL is safe
				tr := &http.Transport{
					Proxy:           http.ProxyFromEnvironment,
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				}
				client := &http.Client{Transport: tr}

				tokenResponse, err := client.Do(tokenReq)
				o.Expect(err).NotTo(o.HaveOccurred())
				response, err := httputil.DumpResponse(tokenResponse, true)
				o.Expect(err).NotTo(o.HaveOccurred())
				framework.Logf("%v", string(response))
				o.Expect(tokenResponse.StatusCode).To(o.Equal(http.StatusOK))
			}
		})
	})
})

func queryClientAuthRequest(host string) *http.Request {
	req, err := http.NewRequest("POST", "https://"+host+"/oauth/token?"+
		"grant_type=authorization_code&"+
		"code="+authzTokenName+"&"+
		"client_id=oauth-client-with-plus&"+
		"client_secret=secret%2Bwith%2Bplus",
		nil)
	o.Expect(err).NotTo(o.HaveOccurred())

	return req
}

func bodyClientAuthRequest(host string) *http.Request {
	req, err := http.NewRequest("POST", "https://"+host+"/oauth/token",
		bytes.NewBufferString("grant_type=authorization_code&"+
			"code="+authzTokenName+"&"+
			"client_id=oauth-client-with-plus&"+
			"client_secret=secret%2Bwith%2Bplus",
		),
	)
	o.Expect(err).NotTo(o.HaveOccurred())

	return req
}

func headerClientAuthRequest(host string) *http.Request {
	req, err := http.NewRequest("POST", "https://"+host+"/oauth/token",
		bytes.NewBufferString("grant_type=authorization_code&"+
			"code="+authzTokenName,
		),
	)
	o.Expect(err).NotTo(o.HaveOccurred())

	req.Header.Set("Authorization", fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte("oauth-client-with-plus:secret+with+plus"))))
	return req
}
<|endoftext|>


<commit_before>package gpio

import (
	"strings"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*MotorDriver)(nil)

func initTestMotorDriver() *MotorDriver {
	return NewMotorDriver(newGpioTestAdaptor(), "1")
}

func TestMotorDriver(t *testing.T) {
	d := NewMotorDriver(newGpioTestAdaptor(), "1")
	gobottest.Refute(t, d.Connection(), nil)

}
func TestMotorDriverStart(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, d.Start(), nil)
}

func TestMotorDriverHalt(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, d.Halt(), nil)
}

func TestMotorDriverIsOn(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.CurrentState = 1
	gobottest.Assert(t, d.IsOn(), true)
	d.CurrentMode = "analog"
	d.CurrentSpeed = 100
	gobottest.Assert(t, d.IsOn(), true)
}

func TestMotorDriverIsOff(t *testing.T) {
	d := initTestMotorDriver()
	d.Off()
	gobottest.Assert(t, d.IsOff(), true)
}

func TestMotorDriverOn(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.On()
	gobottest.Assert(t, d.CurrentState, uint8(1))
	d.CurrentMode = "analog"
	d.CurrentSpeed = 0
	d.On()
	gobottest.Assert(t, d.CurrentSpeed, uint8(255))
}

func TestMotorDriverOff(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.Off()
	gobottest.Assert(t, d.CurrentState, uint8(0))
	d.CurrentMode = "analog"
	d.CurrentSpeed = 100
	d.Off()
	gobottest.Assert(t, d.CurrentSpeed, uint8(0))
}

func TestMotorDriverToggle(t *testing.T) {
	d := initTestMotorDriver()
	d.Off()
	d.Toggle()
	gobottest.Assert(t, d.IsOn(), true)
	d.Toggle()
	gobottest.Assert(t, d.IsOn(), false)
}

func TestMotorDriverMin(t *testing.T) {
	d := initTestMotorDriver()
	d.Min()
}

func TestMotorDriverMax(t *testing.T) {
	d := initTestMotorDriver()
	d.Max()
}

func TestMotorDriverSpeed(t *testing.T) {
	d := initTestMotorDriver()
	d.Speed(100)
}

func TestMotorDriverForward(t *testing.T) {
	d := initTestMotorDriver()
	d.Forward(100)
	gobottest.Assert(t, d.CurrentSpeed, uint8(100))
	gobottest.Assert(t, d.CurrentDirection, "forward")
}

func TestMotorDriverBackward(t *testing.T) {
	d := initTestMotorDriver()
	d.Backward(100)
	gobottest.Assert(t, d.CurrentSpeed, uint8(100))
	gobottest.Assert(t, d.CurrentDirection, "backward")
}

func TestMotorDriverDirection(t *testing.T) {
	d := initTestMotorDriver()
	d.Direction("none")
	d.DirectionPin = "2"
	d.Direction("forward")
	d.Direction("backward")
}

func TestMotorDriverDigital(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.ForwardPin = "2"
	d.BackwardPin = "3"

	d.On()
	gobottest.Assert(t, d.CurrentState, uint8(1))
	d.Off()
	gobottest.Assert(t, d.CurrentState, uint8(0))
}

func TestMotorDriverDefaultName(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, strings.HasPrefix(d.Name(), "Motor"), true)
}

func TestMotorDriverSetName(t *testing.T) {
	d := initTestMotorDriver()
	d.SetName("mybot")
	gobottest.Assert(t, d.Name(), "mybot")
}
<commit_msg>Fix test on mode change when speed is set<commit_after>package gpio

import (
	"strings"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*MotorDriver)(nil)

func initTestMotorDriver() *MotorDriver {
	return NewMotorDriver(newGpioTestAdaptor(), "1")
}

func TestMotorDriver(t *testing.T) {
	d := NewMotorDriver(newGpioTestAdaptor(), "1")
	gobottest.Refute(t, d.Connection(), nil)

}
func TestMotorDriverStart(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, d.Start(), nil)
}

func TestMotorDriverHalt(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, d.Halt(), nil)
}

func TestMotorDriverIsOn(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.CurrentState = 1
	gobottest.Assert(t, d.IsOn(), true)
	d.CurrentMode = "analog"
	d.CurrentSpeed = 100
	gobottest.Assert(t, d.IsOn(), true)
}

func TestMotorDriverIsOff(t *testing.T) {
	d := initTestMotorDriver()
	d.Off()
	gobottest.Assert(t, d.IsOff(), true)
}

func TestMotorDriverOn(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.On()
	gobottest.Assert(t, d.CurrentState, uint8(1))
	d.CurrentMode = "analog"
	d.CurrentSpeed = 0
	d.On()
	gobottest.Assert(t, d.CurrentSpeed, uint8(255))
}

func TestMotorDriverOff(t *testing.T) {
	d := initTestMotorDriver()
	d.CurrentMode = "digital"
	d.Off()
	gobottest.Assert(t, d.CurrentState, uint8(0))
	d.CurrentMode = "analog"
	d.CurrentSpeed = 100
	d.Off()
	gobottest.Assert(t, d.CurrentSpeed, uint8(0))
}

func TestMotorDriverToggle(t *testing.T) {
	d := initTestMotorDriver()
	d.Off()
	d.Toggle()
	gobottest.Assert(t, d.IsOn(), true)
	d.Toggle()
	gobottest.Assert(t, d.IsOn(), false)
}

func TestMotorDriverMin(t *testing.T) {
	d := initTestMotorDriver()
	d.Min()
}

func TestMotorDriverMax(t *testing.T) {
	d := initTestMotorDriver()
	d.Max()
}

func TestMotorDriverSpeed(t *testing.T) {
	d := initTestMotorDriver()
	d.Speed(100)
}

func TestMotorDriverForward(t *testing.T) {
	d := initTestMotorDriver()
	d.Forward(100)
	gobottest.Assert(t, d.CurrentSpeed, uint8(100))
	gobottest.Assert(t, d.CurrentDirection, "forward")
}

func TestMotorDriverBackward(t *testing.T) {
	d := initTestMotorDriver()
	d.Backward(100)
	gobottest.Assert(t, d.CurrentSpeed, uint8(100))
	gobottest.Assert(t, d.CurrentDirection, "backward")
}

func TestMotorDriverDirection(t *testing.T) {
	d := initTestMotorDriver()
	d.Direction("none")
	d.DirectionPin = "2"
	d.Direction("forward")
	d.Direction("backward")
}

func TestMotorDriverDigital(t *testing.T) {
	d := initTestMotorDriver()
	d.SpeedPin = "" // Disable speed
	d.CurrentMode = "digital"
	d.ForwardPin = "2"
	d.BackwardPin = "3"

	d.On()
	gobottest.Assert(t, d.CurrentState, uint8(1))
	d.Off()
	gobottest.Assert(t, d.CurrentState, uint8(0))
}

func TestMotorDriverDefaultName(t *testing.T) {
	d := initTestMotorDriver()
	gobottest.Assert(t, strings.HasPrefix(d.Name(), "Motor"), true)
}

func TestMotorDriverSetName(t *testing.T) {
	d := initTestMotorDriver()
	d.SetName("mybot")
	gobottest.Assert(t, d.Name(), "mybot")
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Workiva, LLC

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

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

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

package set

import (
	"reflect"
	"strconv"
	"testing"
)

func TestAddDuplicateItem(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Add(`test`)

	if !reflect.DeepEqual([]interface{}{`test`}, set.Flatten()) {
		t.Errorf(`Incorrect result returned: %+v`, set.Flatten())
	}
}

func TestAddItems(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Add(`test1`)

	firstSeen := false
	secondSeen := false
	// order is not guaranteed
	for _, item := range set.Flatten() {
		if item.(string) == `test` {
			firstSeen = true
		} else if item.(string) == `test1` {
			secondSeen = true
		}
	}

	if !firstSeen || !secondSeen {
		t.Errorf(`Not all items seen in set.`)
	}
}

func TestRemove(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Remove(`test`)

	if !reflect.DeepEqual([]interface{}{}, set.Flatten()) {
		t.Errorf(`Incorrect result returned: %+v`, set.Flatten())
	}
}

func TestExists(t *testing.T) {
	set := New()
	set.Add(`test`)

	if !set.Exists(`test`) {
		t.Errorf(`Correct existence not determined`)
	}

	if set.Exists(`test1`) {
		t.Errorf(`Correct nonexistence not determined.`)
	}
}

func TestLen(t *testing.T) {
	set := New()
	set.Add(`test`)

	if set.Len() != 1 {
		t.Errorf(`Expected len: %d, received: %d`, 1, set.Len())
	}

	set.Add(`test1`)
	if set.Len() != 2 {
		t.Errorf(`Expected len: %d, received: %d`, 2, set.Len())
	}
}

func TestFlattenCaches(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)

	set.Flatten()

	if len(set.flattened) != 1 {
		t.Errorf(`Expected len: %d, received: %d`, 1, len(set.flattened))
	}
}

func TestAddClearsCache(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)
	set.Flatten()

	set.Add(item)

	if len(set.flattened) != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, len(set.flattened))
	}

	item = `test2`
	set.Add(item)

	if set.flattened != nil {
		t.Errorf(`Cache not cleared.`)
	}
}

func TestDeleteClearsCache(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)
	set.Flatten()

	set.Remove(item)

	if set.flattened != nil {
		t.Errorf(`Cache not cleared.`)
	}
}

func TestAll(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)

	result := set.All(item)
	if !result {
		t.Errorf(`Expected true.`)
	}

	itemTwo := `test1`

	result = set.All(item, itemTwo)
	if result {
		t.Errorf(`Expected false.`)
	}
}

func TestClear(t *testing.T) {
	set := New()
	set.Add(`test`)

	set.Clear()

	if set.Len() != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, set.Len())
	}
}

func BenchmarkFlatten(b *testing.B) {
	set := New()
	for i := 0; i < 50; i++ {
		item := strconv.Itoa(i)
		set.Add(item)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Flatten()
	}
}

func BenchmarkLen(b *testing.B) {
	set := New()
	for i := 0; i < 50; i++ {
		item := strconv.Itoa(i)
		set.Add(item)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Len()
	}
}

func BenchmarkExists(b *testing.B) {
	set := New()
	set.Add(1)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Exists(1)
	}
}

func BenchmarkClear(b *testing.B) {
	for i := 0; i < b.N; i++ {
		set.Clear()
	}
}
<commit_msg>[Set] Fix BenchmarkClear again<commit_after>/*
Copyright 2014 Workiva, LLC

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

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

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

package set

import (
	"reflect"
	"strconv"
	"testing"
)

func TestAddDuplicateItem(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Add(`test`)

	if !reflect.DeepEqual([]interface{}{`test`}, set.Flatten()) {
		t.Errorf(`Incorrect result returned: %+v`, set.Flatten())
	}
}

func TestAddItems(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Add(`test1`)

	firstSeen := false
	secondSeen := false
	// order is not guaranteed
	for _, item := range set.Flatten() {
		if item.(string) == `test` {
			firstSeen = true
		} else if item.(string) == `test1` {
			secondSeen = true
		}
	}

	if !firstSeen || !secondSeen {
		t.Errorf(`Not all items seen in set.`)
	}
}

func TestRemove(t *testing.T) {
	set := New()
	set.Add(`test`)
	set.Remove(`test`)

	if !reflect.DeepEqual([]interface{}{}, set.Flatten()) {
		t.Errorf(`Incorrect result returned: %+v`, set.Flatten())
	}
}

func TestExists(t *testing.T) {
	set := New()
	set.Add(`test`)

	if !set.Exists(`test`) {
		t.Errorf(`Correct existence not determined`)
	}

	if set.Exists(`test1`) {
		t.Errorf(`Correct nonexistence not determined.`)
	}
}

func TestLen(t *testing.T) {
	set := New()
	set.Add(`test`)

	if set.Len() != 1 {
		t.Errorf(`Expected len: %d, received: %d`, 1, set.Len())
	}

	set.Add(`test1`)
	if set.Len() != 2 {
		t.Errorf(`Expected len: %d, received: %d`, 2, set.Len())
	}
}

func TestFlattenCaches(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)

	set.Flatten()

	if len(set.flattened) != 1 {
		t.Errorf(`Expected len: %d, received: %d`, 1, len(set.flattened))
	}
}

func TestAddClearsCache(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)
	set.Flatten()

	set.Add(item)

	if len(set.flattened) != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, len(set.flattened))
	}

	item = `test2`
	set.Add(item)

	if set.flattened != nil {
		t.Errorf(`Cache not cleared.`)
	}
}

func TestDeleteClearsCache(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)
	set.Flatten()

	set.Remove(item)

	if set.flattened != nil {
		t.Errorf(`Cache not cleared.`)
	}
}

func TestAll(t *testing.T) {
	set := New()
	item := `test`
	set.Add(item)

	result := set.All(item)
	if !result {
		t.Errorf(`Expected true.`)
	}

	itemTwo := `test1`

	result = set.All(item, itemTwo)
	if result {
		t.Errorf(`Expected false.`)
	}
}

func TestClear(t *testing.T) {
	set := New()
	set.Add(`test`)

	set.Clear()

	if set.Len() != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, set.Len())
	}
}

func BenchmarkFlatten(b *testing.B) {
	set := New()
	for i := 0; i < 50; i++ {
		item := strconv.Itoa(i)
		set.Add(item)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Flatten()
	}
}

func BenchmarkLen(b *testing.B) {
	set := New()
	for i := 0; i < 50; i++ {
		item := strconv.Itoa(i)
		set.Add(item)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Len()
	}
}

func BenchmarkExists(b *testing.B) {
	set := New()
	set.Add(1)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set.Exists(1)
	}
}

func BenchmarkClear(b *testing.B) {
	set := New()
	for i := 0; i < b.N; i++ {
		set.Clear()
	}
}
<|endoftext|>


<commit_before>package snappy

import (
	"io"
	"sync"

	"github.com/klauspost/compress/snappy"
)

// Framing is an enumeration type used to enable or disable xerial framing of
// snappy messages.
type Framing int

const (
	Framed Framing = iota
	Unframed
)

var (
	readerPool sync.Pool
	writerPool sync.Pool
)

// Codec is the implementation of a compress.Codec which supports creating
// readers and writers for kafka messages compressed with snappy.
type Codec struct {
	// An optional framing to apply to snappy compression.
	//
	// Default to Framed.
	Framing Framing
}

// Code implements the compress.Codec interface.
func (c *Codec) Code() int8 { return 2 }

// Name implements the compress.Codec interface.
func (c *Codec) Name() string { return "snappy" }

// NewReader implements the compress.Codec interface.
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
	x, _ := readerPool.Get().(*xerialReader)
	if x != nil {
		x.Reset(r)
	} else {
		x = &xerialReader{
			reader: r,
			decode: snappy.Decode,
		}
	}
	return &reader{xerialReader: x}
}

// NewWriter implements the compress.Codec interface.
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
	x, _ := writerPool.Get().(*xerialWriter)
	if x != nil {
		x.Reset(w)
	} else {
		x = &xerialWriter{
			writer: w,
			encode: snappy.Encode,
		}
	}
	x.framed = c.Framing == Framed
	return &writer{xerialWriter: x}
}

type reader struct{ *xerialReader }

func (r *reader) Close() (err error) {
	if x := r.xerialReader; x != nil {
		r.xerialReader = nil
		x.Reset(nil)
		readerPool.Put(x)
	}
	return
}

type writer struct{ *xerialWriter }

func (w *writer) Close() (err error) {
	if x := w.xerialWriter; x != nil {
		w.xerialWriter = nil
		err = x.Flush()
		x.Reset(nil)
		writerPool.Put(x)
	}
	return
}
<commit_msg>snappy: configurable compression level<commit_after>package snappy

import (
	"io"
	"sync"

	"github.com/klauspost/compress/s2"
	"github.com/klauspost/compress/snappy"
)

// Framing is an enumeration type used to enable or disable xerial framing of
// snappy messages.
type Framing int

const (
	Framed Framing = iota
	Unframed
)

// Compression level.
type Compression int

const (
	DefaultCompression Compression = iota
	FasterCompression
	BetterCompression
	BestCompression
)

var (
	readerPool sync.Pool
	writerPool sync.Pool
)

// Codec is the implementation of a compress.Codec which supports creating
// readers and writers for kafka messages compressed with snappy.
type Codec struct {
	// An optional framing to apply to snappy compression.
	//
	// Default to Framed.
	Framing Framing

	// Compression level.
	Compression Compression
}

// Code implements the compress.Codec interface.
func (c *Codec) Code() int8 { return 2 }

// Name implements the compress.Codec interface.
func (c *Codec) Name() string { return "snappy" }

// NewReader implements the compress.Codec interface.
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
	x, _ := readerPool.Get().(*xerialReader)
	if x != nil {
		x.Reset(r)
	} else {
		x = &xerialReader{
			reader: r,
			decode: snappy.Decode,
		}
	}
	return &reader{xerialReader: x}
}

// NewWriter implements the compress.Codec interface.
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
	x, _ := writerPool.Get().(*xerialWriter)
	if x != nil {
		x.Reset(w)
	} else {
		x = &xerialWriter{writer: w}
	}
	x.framed = c.Framing == Framed
	switch c.Compression {
	case FasterCompression:
		x.encode = s2.EncodeSnappy
	case BetterCompression:
		x.encode = s2.EncodeSnappyBetter
	case BestCompression:
		x.encode = s2.EncodeSnappyBest
	default:
		x.encode = snappy.Encode // aka. s2.EncodeSnappyBetter
	}
	return &writer{xerialWriter: x}
}

type reader struct{ *xerialReader }

func (r *reader) Close() (err error) {
	if x := r.xerialReader; x != nil {
		r.xerialReader = nil
		x.Reset(nil)
		readerPool.Put(x)
	}
	return
}

type writer struct{ *xerialWriter }

func (w *writer) Close() (err error) {
	if x := w.xerialWriter; x != nil {
		w.xerialWriter = nil
		err = x.Flush()
		x.Reset(nil)
		writerPool.Put(x)
	}
	return
}
<|endoftext|>


<commit_before>package models

type PermissionView struct {
	//Username         string `json:"username,omitempty"`
	Email            string `json:"email,omitempty"`
	ServerIdentifier string `json:"serverIdentifier,omitempty"`

	EditServerData    bool `json:"editServerData,omitempty,string"`
	EditServerUsers   bool `json:"editServerUsers,omitempty,string"`
	InstallServer     bool `json:"installServer,omitempty,string"`
	UpdateServer      bool `json:"-"` //this is unused currently
	ViewServerConsole bool `json:"viewServerConsole,omitempty,string"`
	SendServerConsole bool `json:"sendServerConsole,omitempty,string"`
	StopServer        bool `json:"stopServer,omitempty,string"`
	StartServer       bool `json:"startServer,omitempty,string"`
	ViewServerStats   bool `json:"viewServerStats,omitempty,string"`
	ViewServerFiles   bool `json:"viewServerFiles,omitempty,string"`
	SFTPServer        bool `json:"sftpServer,omitempty,string"`
	PutServerFiles    bool `json:"putServerFiles,omitempty,string"`

	Admin           bool `json:"admin,omitempty,string"`
	ViewServer      bool `json:"viewServers,omitempty,string"`
	CreateServer    bool `json:"createServers,omitempty,string"`
	ViewNodes       bool `json:"viewNodes,omitempty,string"`
	EditNodes       bool `json:"editNodes,omitempty,string"`
	DeployNodes     bool `json:"deployNodes,omitempty,string"`
	ViewTemplates   bool `json:"viewTemplates,omitempty,string"`
	EditUsers       bool `json:"editUsers,omitempty,string"`
	ViewUsers       bool `json:"viewUsers,omitempty,string"`
	EditServerAdmin bool `json:"editServerAdmin,omitempty,string"`
	DeleteServer    bool `json:"deleteServers,omitempty,string"`
}

func FromPermission(p *Permissions) *PermissionView {
	model := &PermissionView{
		//Username: p.User.Username,
		Email: p.User.Email,
	}

	//only show server specific perms
	if p.ServerIdentifier != nil {
		model.ServerIdentifier = *p.ServerIdentifier
		model.EditServerData = p.EditServerData
		model.EditServerUsers = p.EditServerUsers
		model.InstallServer = p.InstallServer
		model.UpdateServer = p.UpdateServer
		model.ViewServerConsole = p.ViewServerConsole
		model.SendServerConsole = p.SendServerConsole
		model.StopServer = p.StopServer
		model.StartServer = p.StartServer
		model.ViewServerStats = p.ViewServerStats
		model.ViewServerFiles = p.ViewServerFiles
		model.SFTPServer = p.SFTPServer
		model.PutServerFiles = p.PutServerFiles
	} else {
		model.Admin = p.Admin
		model.ViewServer = p.ViewServer
		model.CreateServer = p.CreateServer
		model.ViewNodes = p.ViewNodes
		model.EditNodes = p.EditNodes
		model.DeployNodes = p.DeployNodes
		model.ViewTemplates = p.ViewTemplates
		model.EditUsers = p.EditUsers
		model.ViewUsers = p.ViewUsers
		model.EditServerAdmin = p.EditServerAdmin
		model.DeleteServer = p.DeleteServer
	}

	return model
}

//Copies perms from the view to the model
//This will only copy what it knows about the server
func (p *PermissionView) CopyTo(model *Permissions, copyAdminFlags bool) {
	if model.ServerIdentifier != nil {
		model.EditServerData = p.EditServerData
		model.EditServerUsers = p.EditServerUsers
		model.InstallServer = p.InstallServer
		model.UpdateServer = p.UpdateServer
		model.ViewServerConsole = p.ViewServerConsole
		model.SendServerConsole = p.SendServerConsole
		model.StopServer = p.StopServer
		model.StartServer = p.StartServer
		model.ViewServerStats = p.ViewServerStats
		model.ViewServerFiles = p.ViewServerFiles
		model.SFTPServer = p.SFTPServer
		model.PutServerFiles = p.PutServerFiles
	} else if copyAdminFlags {
		model.Admin = p.Admin
		model.ViewServer = p.ViewServer
		model.CreateServer = p.CreateServer
		model.ViewNodes = p.ViewNodes
		model.EditNodes = p.EditNodes
		model.DeployNodes = p.DeployNodes
		model.ViewTemplates = p.ViewTemplates
		model.EditUsers = p.EditUsers
		model.ViewUsers = p.ViewUsers
		model.EditServerAdmin = p.EditServerAdmin
		model.DeleteServer = p.DeleteServer
	}
}
<commit_msg>Re-expose username in returned results<commit_after>package models

type PermissionView struct {
	Username         string `json:"username,omitempty"`
	Email            string `json:"email,omitempty"`
	ServerIdentifier string `json:"serverIdentifier,omitempty"`

	EditServerData    bool `json:"editServerData,omitempty,string"`
	EditServerUsers   bool `json:"editServerUsers,omitempty,string"`
	InstallServer     bool `json:"installServer,omitempty,string"`
	UpdateServer      bool `json:"-"` //this is unused currently
	ViewServerConsole bool `json:"viewServerConsole,omitempty,string"`
	SendServerConsole bool `json:"sendServerConsole,omitempty,string"`
	StopServer        bool `json:"stopServer,omitempty,string"`
	StartServer       bool `json:"startServer,omitempty,string"`
	ViewServerStats   bool `json:"viewServerStats,omitempty,string"`
	ViewServerFiles   bool `json:"viewServerFiles,omitempty,string"`
	SFTPServer        bool `json:"sftpServer,omitempty,string"`
	PutServerFiles    bool `json:"putServerFiles,omitempty,string"`

	Admin           bool `json:"admin,omitempty,string"`
	ViewServer      bool `json:"viewServers,omitempty,string"`
	CreateServer    bool `json:"createServers,omitempty,string"`
	ViewNodes       bool `json:"viewNodes,omitempty,string"`
	EditNodes       bool `json:"editNodes,omitempty,string"`
	DeployNodes     bool `json:"deployNodes,omitempty,string"`
	ViewTemplates   bool `json:"viewTemplates,omitempty,string"`
	EditUsers       bool `json:"editUsers,omitempty,string"`
	ViewUsers       bool `json:"viewUsers,omitempty,string"`
	EditServerAdmin bool `json:"editServerAdmin,omitempty,string"`
	DeleteServer    bool `json:"deleteServers,omitempty,string"`
}

func FromPermission(p *Permissions) *PermissionView {
	model := &PermissionView{
		Username: p.User.Username,
		Email: p.User.Email,
	}

	//only show server specific perms
	if p.ServerIdentifier != nil {
		model.ServerIdentifier = *p.ServerIdentifier
		model.EditServerData = p.EditServerData
		model.EditServerUsers = p.EditServerUsers
		model.InstallServer = p.InstallServer
		model.UpdateServer = p.UpdateServer
		model.ViewServerConsole = p.ViewServerConsole
		model.SendServerConsole = p.SendServerConsole
		model.StopServer = p.StopServer
		model.StartServer = p.StartServer
		model.ViewServerStats = p.ViewServerStats
		model.ViewServerFiles = p.ViewServerFiles
		model.SFTPServer = p.SFTPServer
		model.PutServerFiles = p.PutServerFiles
	} else {
		model.Admin = p.Admin
		model.ViewServer = p.ViewServer
		model.CreateServer = p.CreateServer
		model.ViewNodes = p.ViewNodes
		model.EditNodes = p.EditNodes
		model.DeployNodes = p.DeployNodes
		model.ViewTemplates = p.ViewTemplates
		model.EditUsers = p.EditUsers
		model.ViewUsers = p.ViewUsers
		model.EditServerAdmin = p.EditServerAdmin
		model.DeleteServer = p.DeleteServer
	}

	return model
}

//Copies perms from the view to the model
//This will only copy what it knows about the server
func (p *PermissionView) CopyTo(model *Permissions, copyAdminFlags bool) {
	if model.ServerIdentifier != nil {
		model.EditServerData = p.EditServerData
		model.EditServerUsers = p.EditServerUsers
		model.InstallServer = p.InstallServer
		model.UpdateServer = p.UpdateServer
		model.ViewServerConsole = p.ViewServerConsole
		model.SendServerConsole = p.SendServerConsole
		model.StopServer = p.StopServer
		model.StartServer = p.StartServer
		model.ViewServerStats = p.ViewServerStats
		model.ViewServerFiles = p.ViewServerFiles
		model.SFTPServer = p.SFTPServer
		model.PutServerFiles = p.PutServerFiles
	} else if copyAdminFlags {
		model.Admin = p.Admin
		model.ViewServer = p.ViewServer
		model.CreateServer = p.CreateServer
		model.ViewNodes = p.ViewNodes
		model.EditNodes = p.EditNodes
		model.DeployNodes = p.DeployNodes
		model.ViewTemplates = p.ViewTemplates
		model.EditUsers = p.EditUsers
		model.ViewUsers = p.ViewUsers
		model.EditServerAdmin = p.EditServerAdmin
		model.DeleteServer = p.DeleteServer
	}
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/golang/glog"
)

var (
	nvidiaSmiPath = flag.String("nvidia-smi-path", "/usr/local/nvidia/bin/nvidia-smi", "Path where nvidia-smi is installed.")
	gpuConfigFile = flag.String("gpu-config", "/etc/nvidia/gpu_config.json", "File with GPU configurations for device plugin")
)

var partitionSizeToProfileID = map[string]string{
	"1g.5gb":  "19",
	"2g.10gb": "14",
	"3g.20gb": "9",
	"4g.20gb": "5",
	"7g.40gb": "0",
}

var partitionSizeMaxCount = map[string]int{
	"1g.5gb":  7,
	"2g.10gb": 3,
	"3g.20gb": 2,
	"4g.20gb": 1,
	"7g.40gb": 1,
}

// GPUConfig stores the settings used to configure the GPUs on a node.
type GPUConfig struct {
	GPUPartitionSize string
}

func main() {
	flag.Parse()

	if _, err := os.Stat(*gpuConfigFile); os.IsNotExist(err) {
		glog.Infof("No GPU config file given, nothing to do.")
		return
	}
	gpuConfig, err := parseGPUConfig(*gpuConfigFile)
	if err != nil {
		glog.Infof("failed to parse GPU config file, taking no action.")
		return
	}
	glog.Infof("Using gpu config: %v", gpuConfig)
	if gpuConfig.GPUPartitionSize == "" {
		glog.Infof("No GPU partitions are required, exiting")
		return
	}

	if _, err := os.Stat(*nvidiaSmiPath); os.IsNotExist(err) {
		glog.Errorf("nvidia-smi path %s not found: %v", *nvidiaSmiPath, err)
		os.Exit(1)
	}

	migModeEnabled, err := currentMigMode()
	if err != nil {
		glog.Errorf("Failed to check if MIG mode is enabled: %v", err)
		os.Exit(1)
	}
	if !migModeEnabled {
		glog.Infof("MIG mode is not enabled. Enabling now.")
		if err := enableMigMode(); err != nil {
			glog.Errorf("Failed to enable MIG mode: %v", err)
			os.Exit(1)
		}
		glog.Infof("Rebooting node to enable MIG mode")
		if err := rebootNode(); err != nil {
			glog.Errorf("Failed to trigger node reboot after enabling MIG mode: %v", err)
		}

		// Exit, since we cannot proceed until node has rebooted, for MIG changes to take effect.
		os.Exit(1)
	}

	glog.Infof("MIG mode is enabled on all GPUs, proceeding to create GPU partitions.")

	glog.Infof("Cleaning up any existing GPU partitions")
	if err := cleanupAllGPUPartitions(); err != nil {
		glog.Errorf("Failed to cleanup GPU partitions: %v", err)
		os.Exit(1)
	}

	glog.Infof("Creating new GPU partitions")
	if err := createGPUPartitions(gpuConfig.GPUPartitionSize); err != nil {
		glog.Errorf("Failed to create GPU partitions: %v", err)
		os.Exit(1)
	}

	glog.Infof("Running %s", *nvidiaSmiPath)
	out, err := exec.Command(*nvidiaSmiPath).Output()
	if err != nil {
		glog.Errorf("Failed to run nvidia-smi, output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

}

func parseGPUConfig(gpuConfigFile string) (GPUConfig, error) {
	var gpuConfig GPUConfig

	gpuConfigContent, err := ioutil.ReadFile(gpuConfigFile)
	if err != nil {
		return gpuConfig, fmt.Errorf("unable to read gpu config file %s: %v", gpuConfigFile, err)
	}

	if err = json.Unmarshal(gpuConfigContent, &gpuConfig); err != nil {
		return gpuConfig, fmt.Errorf("failed to parse GPU config file contents: %s, error: %v", gpuConfigContent, err)
	}
	return gpuConfig, nil
}

// currentMigMode returns whether mig mode is currently enabled all GPUs attached to this node.
func currentMigMode() (bool, error) {
	out, err := exec.Command(*nvidiaSmiPath, "--query-gpu=mig.mode.current", "--format=csv,noheader").Output()
	if err != nil {
		return false, err
	}
	if strings.HasPrefix(string(out), "Enabled") {
		return true, nil
	}
	if strings.HasPrefix(string(out), "Disabled") {
		return false, nil
	}
	return false, fmt.Errorf("nvidia-smi returned invalid output: %s", out)
}

// enableMigMode enables MIG mode on all GPUs attached to the node. Requires node restart to take effect.
func enableMigMode() error {
	return exec.Command(*nvidiaSmiPath, "-mig", "1").Run()
}

func rebootNode() error {
	return ioutil.WriteFile("/proc/sysrq-trigger", []byte("b"), 0644)
}

func cleanupAllGPUPartitions() error {
	args := []string{"mig", "-dci"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err := exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil && !strings.Contains(string(out), "No GPU instances found") {
		return fmt.Errorf("failed to destroy compute instance, nvidia-smi output: %s, error: %v ", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	args = []string{"mig", "-dgi"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err = exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil && !strings.Contains(string(out), "No GPU instances found") {
		return fmt.Errorf("failed to destroy gpu instance, nvidia-smi output: %s, error: %v ", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))
	return nil
}

func createGPUPartitions(partitionSize string) error {
	p, err := buildPartitionStr(partitionSize)
	if err != nil {
		return err
	}

	args := []string{"mig", "-cgi", p}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err := exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil {
		return fmt.Errorf("failed to create GPU Instances: output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	args = []string{"mig", "-cci"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err = exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil {
		return fmt.Errorf("failed to create compute instances: output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	return nil

}

func buildPartitionStr(partitionSize string) (string, error) {
	if partitionSize == "" {
		return "", nil
	}

	p, ok := partitionSizeToProfileID[partitionSize]
	if !ok {
		return "", fmt.Errorf("%s is not a valid partition size", partitionSize)
	}

	partitionStr := p
	for i := 1; i < partitionSizeMaxCount[partitionSize]; i++ {
		partitionStr += fmt.Sprintf(",%s", p)
	}

	return partitionStr, nil
}
<commit_msg>Perform more graceful node reboot after enabling MIG mode on GPUs<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/golang/glog"
)

var (
	nvidiaSmiPath = flag.String("nvidia-smi-path", "/usr/local/nvidia/bin/nvidia-smi", "Path where nvidia-smi is installed.")
	gpuConfigFile = flag.String("gpu-config", "/etc/nvidia/gpu_config.json", "File with GPU configurations for device plugin")
)

var partitionSizeToProfileID = map[string]string{
	"1g.5gb":  "19",
	"2g.10gb": "14",
	"3g.20gb": "9",
	"4g.20gb": "5",
	"7g.40gb": "0",
}

var partitionSizeMaxCount = map[string]int{
	"1g.5gb":  7,
	"2g.10gb": 3,
	"3g.20gb": 2,
	"4g.20gb": 1,
	"7g.40gb": 1,
}

const SIGRTMIN = 34

// GPUConfig stores the settings used to configure the GPUs on a node.
type GPUConfig struct {
	GPUPartitionSize string
}

func main() {
	flag.Parse()

	if _, err := os.Stat(*gpuConfigFile); os.IsNotExist(err) {
		glog.Infof("No GPU config file given, nothing to do.")
		return
	}
	gpuConfig, err := parseGPUConfig(*gpuConfigFile)
	if err != nil {
		glog.Infof("failed to parse GPU config file, taking no action.")
		return
	}
	glog.Infof("Using gpu config: %v", gpuConfig)
	if gpuConfig.GPUPartitionSize == "" {
		glog.Infof("No GPU partitions are required, exiting")
		return
	}

	if _, err := os.Stat(*nvidiaSmiPath); os.IsNotExist(err) {
		glog.Errorf("nvidia-smi path %s not found: %v", *nvidiaSmiPath, err)
		os.Exit(1)
	}

	migModeEnabled, err := currentMigMode()
	if err != nil {
		glog.Errorf("Failed to check if MIG mode is enabled: %v", err)
		os.Exit(1)
	}
	if !migModeEnabled {
		glog.Infof("MIG mode is not enabled. Enabling now.")
		if err := enableMigMode(); err != nil {
			glog.Errorf("Failed to enable MIG mode: %v", err)
			os.Exit(1)
		}
		glog.Infof("Rebooting node to enable MIG mode")
		if err := rebootNode(); err != nil {
			glog.Errorf("Failed to trigger node reboot after enabling MIG mode: %v", err)
		}

		// Exit, since we cannot proceed until node has rebooted, for MIG changes to take effect.
		os.Exit(1)
	}

	glog.Infof("MIG mode is enabled on all GPUs, proceeding to create GPU partitions.")

	glog.Infof("Cleaning up any existing GPU partitions")
	if err := cleanupAllGPUPartitions(); err != nil {
		glog.Errorf("Failed to cleanup GPU partitions: %v", err)
		os.Exit(1)
	}

	glog.Infof("Creating new GPU partitions")
	if err := createGPUPartitions(gpuConfig.GPUPartitionSize); err != nil {
		glog.Errorf("Failed to create GPU partitions: %v", err)
		os.Exit(1)
	}

	glog.Infof("Running %s", *nvidiaSmiPath)
	out, err := exec.Command(*nvidiaSmiPath).Output()
	if err != nil {
		glog.Errorf("Failed to run nvidia-smi, output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

}

func parseGPUConfig(gpuConfigFile string) (GPUConfig, error) {
	var gpuConfig GPUConfig

	gpuConfigContent, err := ioutil.ReadFile(gpuConfigFile)
	if err != nil {
		return gpuConfig, fmt.Errorf("unable to read gpu config file %s: %v", gpuConfigFile, err)
	}

	if err = json.Unmarshal(gpuConfigContent, &gpuConfig); err != nil {
		return gpuConfig, fmt.Errorf("failed to parse GPU config file contents: %s, error: %v", gpuConfigContent, err)
	}
	return gpuConfig, nil
}

// currentMigMode returns whether mig mode is currently enabled all GPUs attached to this node.
func currentMigMode() (bool, error) {
	out, err := exec.Command(*nvidiaSmiPath, "--query-gpu=mig.mode.current", "--format=csv,noheader").Output()
	if err != nil {
		return false, err
	}
	if strings.HasPrefix(string(out), "Enabled") {
		return true, nil
	}
	if strings.HasPrefix(string(out), "Disabled") {
		return false, nil
	}
	return false, fmt.Errorf("nvidia-smi returned invalid output: %s", out)
}

// enableMigMode enables MIG mode on all GPUs attached to the node. Requires node restart to take effect.
func enableMigMode() error {
	return exec.Command(*nvidiaSmiPath, "-mig", "1").Run()
}

func rebootNode() error {
	// Gracefully reboot systemd: https://man7.org/linux/man-pages/man1/systemd.1.html#SIGNALS
	return syscall.Kill(1, SIGRTMIN+5)
}

func cleanupAllGPUPartitions() error {
	args := []string{"mig", "-dci"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err := exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil && !strings.Contains(string(out), "No GPU instances found") {
		return fmt.Errorf("failed to destroy compute instance, nvidia-smi output: %s, error: %v ", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	args = []string{"mig", "-dgi"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err = exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil && !strings.Contains(string(out), "No GPU instances found") {
		return fmt.Errorf("failed to destroy gpu instance, nvidia-smi output: %s, error: %v ", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))
	return nil
}

func createGPUPartitions(partitionSize string) error {
	p, err := buildPartitionStr(partitionSize)
	if err != nil {
		return err
	}

	args := []string{"mig", "-cgi", p}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err := exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil {
		return fmt.Errorf("failed to create GPU Instances: output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	args = []string{"mig", "-cci"}
	glog.Infof("Running %s %s", *nvidiaSmiPath, strings.Join(args, " "))
	out, err = exec.Command(*nvidiaSmiPath, args...).Output()
	if err != nil {
		return fmt.Errorf("failed to create compute instances: output: %s, error: %v", string(out), err)
	}
	glog.Infof("Output:\n %s", string(out))

	return nil

}

func buildPartitionStr(partitionSize string) (string, error) {
	if partitionSize == "" {
		return "", nil
	}

	p, ok := partitionSizeToProfileID[partitionSize]
	if !ok {
		return "", fmt.Errorf("%s is not a valid partition size", partitionSize)
	}

	partitionStr := p
	for i := 1; i < partitionSizeMaxCount[partitionSize]; i++ {
		partitionStr += fmt.Sprintf(",%s", p)
	}

	return partitionStr, nil
}
<|endoftext|>


<commit_before>package topology

import (
	"k8s.io/kubernetes/pkg/util/json"
	"fmt"
	"regexp"
)

type Context struct {
	inputsDef []Input
	inputs    map[string]interface{}
}

func NewContext(def []Input, inputs map[string]interface{}) *Context {
	if len(def) > 0 && inputs != nil {
		buf, _ := json.Marshal(inputs)
		if buf != nil {
			fmt.Println(string(buf))
		}
	}
	return &Context{inputsDef:def, inputs:inputs}
}

func (p *Context) GetValue(name string) interface{} {
	reg := regexp.MustCompile(`^\$\{(.+)\}$`)
	match := reg.FindStringSubmatch(name)
	if len(match) < 1 {
		return nil
	}
	key := match[1]
	//fmt.Printf("GetValue %s\n", key)
	if key == "" {
		return nil
	} else {
		return p.inputs[key]
	}
}
<commit_msg>fix json<commit_after>package topology

import (
	"fmt"
	"regexp"
	"encoding/json"
)

type Context struct {
	inputsDef []Input
	inputs    map[string]interface{}
}

func NewContext(def []Input, inputs map[string]interface{}) *Context {
	if len(def) > 0 && inputs != nil {
		buf, _ := json.Marshal(inputs)
		if buf != nil {
			fmt.Println(string(buf))
		}
	}
	return &Context{inputsDef:def, inputs:inputs}
}

func (p *Context) GetValue(name string) interface{} {
	reg := regexp.MustCompile(`^\$\{(.+)\}$`)
	match := reg.FindStringSubmatch(name)
	if len(match) < 1 {
		return nil
	}
	key := match[1]
	//fmt.Printf("GetValue %s\n", key)
	if key == "" {
		return nil
	} else {
		return p.inputs[key]
	}
}
<|endoftext|>


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

import (
	"regexp"
	"strings"

	"github.com/gogo/protobuf/proto"

	networking "istio.io/api/networking/v1alpha3"
	"istio.io/istio/pkg/config"
	"istio.io/istio/pkg/config/labels"
	"istio.io/istio/pkg/config/xds"
)

// EnvoyFilterWrapper is a wrapper for the EnvoyFilter api object with pre-processed data
type EnvoyFilterWrapper struct {
	workloadSelector labels.Instance
	Patches          map[networking.EnvoyFilter_ApplyTo][]*EnvoyFilterConfigPatchWrapper
}

// EnvoyFilterConfigPatchWrapper is a wrapper over the EnvoyFilter ConfigPatch api object
// fields are ordered such that this struct is aligned
type EnvoyFilterConfigPatchWrapper struct {
	Value     proto.Message
	Match     *networking.EnvoyFilter_EnvoyConfigObjectMatch
	ApplyTo   networking.EnvoyFilter_ApplyTo
	Operation networking.EnvoyFilter_Patch_Operation
	// Pre-compile the regex from proxy version match in the match
	ProxyVersionRegex *regexp.Regexp
	// ProxyPrefixMatch provides a prefix match for the proxy version. The current API only allows
	// regex match, but as an optimization we can reduce this to a prefix match for common cases.
	// If this is set, ProxyVersionRegex is ignored.
	ProxyPrefixMatch string
}

// wellKnownVersions defines a mapping of well known regex matches to prefix matches
// This is done only as an optimization; behavior should remain the same
// All versions specified by the default installation (Telemetry V2) should be added here.
var wellKnownVersions = map[string]string{
	`^1\.4.*`: "1.4",
	`^1\.5.*`: "1.5",
	`^1\.6.*`: "1.6",
	`^1\.7.*`: "1.7",
	`^1\.8.*`: "1.8",
	// Hopefully we have a better API by 1.9. If not, add it here
}

// convertToEnvoyFilterWrapper converts from EnvoyFilter config to EnvoyFilterWrapper object
func convertToEnvoyFilterWrapper(local *config.Config) *EnvoyFilterWrapper {
	localEnvoyFilter := local.Spec.(*networking.EnvoyFilter)

	out := &EnvoyFilterWrapper{}
	if localEnvoyFilter.WorkloadSelector != nil {
		out.workloadSelector = localEnvoyFilter.WorkloadSelector.Labels
	}
	out.Patches = make(map[networking.EnvoyFilter_ApplyTo][]*EnvoyFilterConfigPatchWrapper)
	for _, cp := range localEnvoyFilter.ConfigPatches {
		cpw := &EnvoyFilterConfigPatchWrapper{
			ApplyTo:   cp.ApplyTo,
			Match:     cp.Match,
			Operation: cp.Patch.Operation,
		}
		var err error
		cpw.Value, err = xds.BuildXDSObjectFromStruct(cp.ApplyTo, cp.Patch.Value)
		// There generally won't be an error here because validation catches mismatched types
		// Should only happen in tests or without validation
		if err != nil {
			log.Errorf("failed to build envoy filter value: %v", err)
		}
		if cp.Match == nil {
			// create a match all object
			cpw.Match = &networking.EnvoyFilter_EnvoyConfigObjectMatch{Context: networking.EnvoyFilter_ANY}
		} else if cp.Match.Proxy != nil && cp.Match.Proxy.ProxyVersion != "" {
			// Attempt to convert regex to a simple prefix match for the common case of matching
			// a standard Istio version. This field should likely be replaced with semver, but for now
			// we can workaround the performance impact of regex
			if prefix, f := wellKnownVersions[cp.Match.Proxy.ProxyVersion]; f {
				cpw.ProxyPrefixMatch = prefix
			} else {
				// pre-compile the regex for proxy version if it exists
				// ignore the error because validation catches invalid regular expressions.
				cpw.ProxyVersionRegex, _ = regexp.Compile(cp.Match.Proxy.ProxyVersion)
			}
		}

		if _, exists := out.Patches[cp.ApplyTo]; !exists {
			out.Patches[cp.ApplyTo] = make([]*EnvoyFilterConfigPatchWrapper, 0)
		}
		if cpw.Operation == networking.EnvoyFilter_Patch_INSERT_AFTER ||
			cpw.Operation == networking.EnvoyFilter_Patch_INSERT_BEFORE ||
			cpw.Operation == networking.EnvoyFilter_Patch_INSERT_FIRST {
			// insert_before, after or first is applicable only for network filter and http filter
			// convert the rest to add
			if cpw.ApplyTo != networking.EnvoyFilter_HTTP_FILTER &&
				cpw.ApplyTo != networking.EnvoyFilter_NETWORK_FILTER &&
				cpw.ApplyTo != networking.EnvoyFilter_HTTP_ROUTE {
				cpw.Operation = networking.EnvoyFilter_Patch_ADD
			}
		}
		out.Patches[cp.ApplyTo] = append(out.Patches[cp.ApplyTo], cpw)
	}
	return out
}

func proxyMatch(proxy *Proxy, cp *EnvoyFilterConfigPatchWrapper) bool {
	if cp.Match.Proxy == nil {
		return true
	}

	if cp.ProxyPrefixMatch != "" {
		if !strings.HasPrefix(proxy.Metadata.IstioVersion, cp.ProxyPrefixMatch) {
			return false
		}
	}
	if cp.ProxyVersionRegex != nil {
		ver := proxy.Metadata.IstioVersion
		if ver == "" {
			// we do not have a proxy version but the user has a regex. so this is a mismatch
			return false
		}
		if !cp.ProxyVersionRegex.MatchString(ver) {
			return false
		}
	}

	for k, v := range cp.Match.Proxy.Metadata {
		if proxy.Metadata.Raw[k] != v {
			return false
		}
	}
	return true
}
<commit_msg>Futureproof telemetry envoyfilters a bit (#28176)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package model

import (
	"regexp"
	"strings"

	"github.com/gogo/protobuf/proto"

	networking "istio.io/api/networking/v1alpha3"
	"istio.io/istio/pkg/config"
	"istio.io/istio/pkg/config/labels"
	"istio.io/istio/pkg/config/xds"
)

// EnvoyFilterWrapper is a wrapper for the EnvoyFilter api object with pre-processed data
type EnvoyFilterWrapper struct {
	workloadSelector labels.Instance
	Patches          map[networking.EnvoyFilter_ApplyTo][]*EnvoyFilterConfigPatchWrapper
}

// EnvoyFilterConfigPatchWrapper is a wrapper over the EnvoyFilter ConfigPatch api object
// fields are ordered such that this struct is aligned
type EnvoyFilterConfigPatchWrapper struct {
	Value     proto.Message
	Match     *networking.EnvoyFilter_EnvoyConfigObjectMatch
	ApplyTo   networking.EnvoyFilter_ApplyTo
	Operation networking.EnvoyFilter_Patch_Operation
	// Pre-compile the regex from proxy version match in the match
	ProxyVersionRegex *regexp.Regexp
	// ProxyPrefixMatch provides a prefix match for the proxy version. The current API only allows
	// regex match, but as an optimization we can reduce this to a prefix match for common cases.
	// If this is set, ProxyVersionRegex is ignored.
	ProxyPrefixMatch string
}

// wellKnownVersions defines a mapping of well known regex matches to prefix matches
// This is done only as an optimization; behavior should remain the same
// All versions specified by the default installation (Telemetry V2) should be added here.
var wellKnownVersions = map[string]string{
	`^1\.4.*`: "1.4",
	`^1\.5.*`: "1.5",
	`^1\.6.*`: "1.6",
	`^1\.7.*`: "1.7",
	`^1\.8.*`: "1.8",
	`^1\.9.*`: "1.9",
	// Hopefully we have a better API by 1.10. If not, add it here
}

// convertToEnvoyFilterWrapper converts from EnvoyFilter config to EnvoyFilterWrapper object
func convertToEnvoyFilterWrapper(local *config.Config) *EnvoyFilterWrapper {
	localEnvoyFilter := local.Spec.(*networking.EnvoyFilter)

	out := &EnvoyFilterWrapper{}
	if localEnvoyFilter.WorkloadSelector != nil {
		out.workloadSelector = localEnvoyFilter.WorkloadSelector.Labels
	}
	out.Patches = make(map[networking.EnvoyFilter_ApplyTo][]*EnvoyFilterConfigPatchWrapper)
	for _, cp := range localEnvoyFilter.ConfigPatches {
		cpw := &EnvoyFilterConfigPatchWrapper{
			ApplyTo:   cp.ApplyTo,
			Match:     cp.Match,
			Operation: cp.Patch.Operation,
		}
		var err error
		cpw.Value, err = xds.BuildXDSObjectFromStruct(cp.ApplyTo, cp.Patch.Value)
		// There generally won't be an error here because validation catches mismatched types
		// Should only happen in tests or without validation
		if err != nil {
			log.Errorf("failed to build envoy filter value: %v", err)
		}
		if cp.Match == nil {
			// create a match all object
			cpw.Match = &networking.EnvoyFilter_EnvoyConfigObjectMatch{Context: networking.EnvoyFilter_ANY}
		} else if cp.Match.Proxy != nil && cp.Match.Proxy.ProxyVersion != "" {
			// Attempt to convert regex to a simple prefix match for the common case of matching
			// a standard Istio version. This field should likely be replaced with semver, but for now
			// we can workaround the performance impact of regex
			if prefix, f := wellKnownVersions[cp.Match.Proxy.ProxyVersion]; f {
				cpw.ProxyPrefixMatch = prefix
			} else {
				// pre-compile the regex for proxy version if it exists
				// ignore the error because validation catches invalid regular expressions.
				cpw.ProxyVersionRegex, _ = regexp.Compile(cp.Match.Proxy.ProxyVersion)
			}
		}

		if _, exists := out.Patches[cp.ApplyTo]; !exists {
			out.Patches[cp.ApplyTo] = make([]*EnvoyFilterConfigPatchWrapper, 0)
		}
		if cpw.Operation == networking.EnvoyFilter_Patch_INSERT_AFTER ||
			cpw.Operation == networking.EnvoyFilter_Patch_INSERT_BEFORE ||
			cpw.Operation == networking.EnvoyFilter_Patch_INSERT_FIRST {
			// insert_before, after or first is applicable only for network filter and http filter
			// convert the rest to add
			if cpw.ApplyTo != networking.EnvoyFilter_HTTP_FILTER &&
				cpw.ApplyTo != networking.EnvoyFilter_NETWORK_FILTER &&
				cpw.ApplyTo != networking.EnvoyFilter_HTTP_ROUTE {
				cpw.Operation = networking.EnvoyFilter_Patch_ADD
			}
		}
		out.Patches[cp.ApplyTo] = append(out.Patches[cp.ApplyTo], cpw)
	}
	return out
}

func proxyMatch(proxy *Proxy, cp *EnvoyFilterConfigPatchWrapper) bool {
	if cp.Match.Proxy == nil {
		return true
	}

	if cp.ProxyPrefixMatch != "" {
		if !strings.HasPrefix(proxy.Metadata.IstioVersion, cp.ProxyPrefixMatch) {
			return false
		}
	}
	if cp.ProxyVersionRegex != nil {
		ver := proxy.Metadata.IstioVersion
		if ver == "" {
			// we do not have a proxy version but the user has a regex. so this is a mismatch
			return false
		}
		if !cp.ProxyVersionRegex.MatchString(ver) {
			return false
		}
	}

	for k, v := range cp.Match.Proxy.Metadata {
		if proxy.Metadata.Raw[k] != v {
			return false
		}
	}
	return true
}
<|endoftext|>


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

package v1beta1

import (
	appsv1 "k8s.io/api/apps/v1"
	policyv1 "k8s.io/api/policy/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

// Constants defining labels
const (
	StatusReady      = "Ready"
	StatusInProgress = "InProgress"
	StatusDisabled   = "Disabled"
)

func (s *ObjectStatus) update(rsrc metav1.Object) {
	ro := rsrc.(runtime.Object)
	gvk := ro.GetObjectKind().GroupVersionKind()
	s.Link = rsrc.GetSelfLink()
	s.Name = rsrc.GetName()
	s.Group = gvk.GroupVersion().String()
	s.Kind = gvk.GroupKind().Kind
	s.Status = StatusReady
}

// ResetComponentList - reset component list objects
func (m *ApplicationStatus) ResetComponentList() {
	m.ComponentList.Objects = []ObjectStatus{}
}

// UpdateStatus the component status
func (m *ApplicationStatus) UpdateStatus(rsrcs []metav1.Object, err error) {
	var ready = true
	for _, r := range rsrcs {
		os := ObjectStatus{}
		os.update(r)
		switch r.(type) {
		case *appsv1.StatefulSet:
			os.Status = stsStatus(r.(*appsv1.StatefulSet))
		case *policyv1.PodDisruptionBudget:
			os.Status = pdbStatus(r.(*policyv1.PodDisruptionBudget))
		}
		m.ComponentList.Objects = append(m.ComponentList.Objects, os)
	}
	for _, os := range m.ComponentList.Objects {
		if os.Status != StatusReady {
			ready = false
		}
	}

	if ready {
		m.Ready("ComponentsReady", "all components ready")
	} else {
		m.NotReady("ComponentsNotReady", "some components not ready")
	}
	if err != nil {
		m.SetCondition(Error, "ErrorSeen", err.Error())
	}
}

// Resource specific logic -----------------------------------

// Statefulset
func stsStatus(rsrc *appsv1.StatefulSet) string {
	if rsrc.Status.ReadyReplicas == *rsrc.Spec.Replicas && rsrc.Status.CurrentReplicas == *rsrc.Spec.Replicas {
		return StatusReady
	}
	return StatusInProgress
}

// PodDisruptionBudget
func pdbStatus(rsrc *policyv1.PodDisruptionBudget) string {
	if rsrc.Status.CurrentHealthy >= rsrc.Status.DesiredHealthy {
		return StatusReady
	}
	return StatusInProgress
}
<commit_msg>Adding ds,rs,deployment<commit_after>/*
Copyright 2018 The Kubernetes Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package v1beta1

import (
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

// Constants defining labels
const (
	StatusReady      = "Ready"
	StatusInProgress = "InProgress"
	StatusDisabled   = "Disabled"
)

func (s *ObjectStatus) update(rsrc metav1.Object) {
	ro := rsrc.(runtime.Object)
	gvk := ro.GetObjectKind().GroupVersionKind()
	s.Link = rsrc.GetSelfLink()
	s.Name = rsrc.GetName()
	s.Group = gvk.GroupVersion().String()
	s.Kind = gvk.GroupKind().Kind
	s.Status = StatusReady
}

// ResetComponentList - reset component list objects
func (m *ApplicationStatus) ResetComponentList() {
	m.ComponentList.Objects = []ObjectStatus{}
}

// UpdateStatus the component status
func (m *ApplicationStatus) UpdateStatus(rsrcs []metav1.Object, err error) {
	var ready = true
	for _, r := range rsrcs {
		os := ObjectStatus{}
		os.update(r)
		switch r.(type) {
		case *appsv1.StatefulSet:
			os.Status = stsStatus(r.(*appsv1.StatefulSet))
		case *policyv1.PodDisruptionBudget:
			os.Status = pdbStatus(r.(*policyv1.PodDisruptionBudget))
		case *appsv1.Deployment:
			os.Status = deploymentStatus(r.(*appsv1.Deployment))
		case *appsv1.ReplicaSet:
			os.Status = replicasetStatus(r.(*appsv1.ReplicaSet))
		case *appsv1.DaemonSet:
			os.Status = daemonsetStatus(r.(*appsv1.DaemonSet))
			//case *corev1.ReplicationController:
			// Pod
			// Ingress
			// Service
			// PersistentVolumeClaim
		}
		m.ComponentList.Objects = append(m.ComponentList.Objects, os)
	}
	for _, os := range m.ComponentList.Objects {
		if os.Status != StatusReady {
			ready = false
		}
	}

	if ready {
		m.Ready("ComponentsReady", "all components ready")
	} else {
		m.NotReady("ComponentsNotReady", "some components not ready")
	}
	if err != nil {
		m.SetCondition(Error, "ErrorSeen", err.Error())
	}
}

// Resource specific logic -----------------------------------

// Statefulset
func stsStatus(rsrc *appsv1.StatefulSet) string {
	if rsrc.Status.ReadyReplicas == *rsrc.Spec.Replicas && rsrc.Status.CurrentReplicas == *rsrc.Spec.Replicas {
		return StatusReady
	}
	return StatusInProgress
}

// Deployment
func deploymentStatus(rsrc *appsv1.Deployment) string {
	status := StatusInProgress
	progress := true
	available := true
	for _, c := range rsrc.Status.Conditions {
		switch c.Type {
		case appsv1.DeploymentProgressing:
			// https://github.com/kubernetes/kubernetes/blob/a3ccea9d8743f2ff82e41b6c2af6dc2c41dc7b10/pkg/controller/deployment/progress.go#L52
			if c.Status != corev1.ConditionTrue || c.Reason != "NewReplicaSetAvailable" {
				progress = false
			}
		case appsv1.DeploymentAvailable:
			if c.Status == corev1.ConditionFalse {
				available = false
			}
		}
	}

	if progress && available {
		status = StatusReady
	}

	return status
}

// Replicaset
func replicasetStatus(rsrc *appsv1.ReplicaSet) string {
	status := StatusInProgress
	failure := false
	for _, c := range rsrc.Status.Conditions {
		switch c.Type {
		// https://github.com/kubernetes/kubernetes/blob/a3ccea9d8743f2ff82e41b6c2af6dc2c41dc7b10/pkg/controller/replicaset/replica_set_utils.go
		case appsv1.ReplicaSetReplicaFailure:
			if c.Status == corev1.ConditionTrue {
				failure = true
			}
		}
	}

	if !failure && rsrc.Status.ReadyReplicas == rsrc.Status.Replicas && rsrc.Status.Replicas == rsrc.Status.AvailableReplicas {
		status = StatusReady
	}

	return status
}

// Daemonset
func daemonsetStatus(rsrc *appsv1.DaemonSet) string {
	status := StatusInProgress
	if rsrc.Status.DesiredNumberScheduled == rsrc.Status.NumberAvailable && rsrc.Status.DesiredNumberScheduled == rsrc.Status.NumberReady {
		status = StatusReady
	}
	return status
}

// PodDisruptionBudget
func pdbStatus(rsrc *policyv1.PodDisruptionBudget) string {
	if rsrc.Status.CurrentHealthy >= rsrc.Status.DesiredHealthy {
		return StatusReady
	}
	return StatusInProgress
}
<|endoftext|>


<commit_before>package templates

import "strings"

func MainHelpTemplate() string {
	return decorate(mainHelpTemplate, false)
}

func MainUsageTemplate() string {
	return decorate(mainUsageTemplate, true)
}

func OptionsHelpTemplate() string {
	return decorate(optionsHelpTemplate, false)
}

func OptionsUsageTemplate() string {
	return decorate(optionsUsageTemplate, false)
}

func decorate(template string, trim bool) string {
	if trim && len(strings.Trim(template, " ")) > 0 {
		template = strings.Trim(template, "\n")
	}
	return template
}

const (
	vars = `{{$isRootCmd := isRootCmd .}}` +
		`{{$rootCmd := rootCmd .}}` +
		`{{$explicitlyExposedFlags := exposed .}}` +
		`{{$localNotPersistentFlags := flagsNotIntersected .LocalFlags .PersistentFlags}}`

	mainHelpTemplate = `{{.Long | trim}}
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`

	mainUsageTemplate = vars + `{{ $cmd := . }}{{ if .HasSubCommands}}
Available Commands: {{range .Commands}}{{if .Runnable}}{{if ne .Name "options"}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}
{{end}}
{{ if or $localNotPersistentFlags.HasFlags $explicitlyExposedFlags.HasFlags}}Options:
{{ if $localNotPersistentFlags.HasFlags}}{{flagsUsages $localNotPersistentFlags}}{{end}}{{ if $explicitlyExposedFlags.HasFlags}}{{flagsUsages $explicitlyExposedFlags}}{{end}}
{{end}}{{ if not $isRootCmd}}Use "{{$rootCmd}} --help" for a list of all commands available in {{$rootCmd}}.
{{end}}{{ if .HasSubCommands }}Use "{{$rootCmd}} <command> --help" for more information about a given command.
{{end}}{{ if and .HasInheritedFlags (not $isRootCmd)}}Use "{{$rootCmd}} options" for a list of global command-line options (applies to all commands).
{{end}}`

	optionsHelpTemplate = ``

	optionsUsageTemplate = `{{ if .HasInheritedFlags}}The following options can be passed to any command:

{{flagsUsages .InheritedFlags}}{{end}}`
)
<commit_msg>Adjust help template to latest version of Cobra<commit_after>package templates

import "strings"

func MainHelpTemplate() string {
	return decorate(mainHelpTemplate, false)
}

func MainUsageTemplate() string {
	return decorate(mainUsageTemplate, true)
}

func OptionsHelpTemplate() string {
	return decorate(optionsHelpTemplate, false)
}

func OptionsUsageTemplate() string {
	return decorate(optionsUsageTemplate, false)
}

func decorate(template string, trim bool) string {
	if trim && len(strings.Trim(template, " ")) > 0 {
		template = strings.Trim(template, "\n")
	}
	return template
}

const (
	vars = `{{$isRootCmd := isRootCmd .}}` +
		`{{$rootCmd := rootCmd .}}` +
		`{{$explicitlyExposedFlags := exposed .}}` +
		`{{$localNotPersistentFlags := flagsNotIntersected .LocalFlags .PersistentFlags}}`

	mainHelpTemplate = `{{.Long | trim}}
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`

	mainUsageTemplate = vars + `{{ $cmd := . }}{{ if .HasRunnableSubCommands}}
Available Commands: {{range .Commands}}{{if and .Runnable (ne .Name "options")}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}
{{end}}
{{ if or $localNotPersistentFlags.HasFlags $explicitlyExposedFlags.HasFlags}}Options:
{{ if $localNotPersistentFlags.HasFlags}}{{flagsUsages $localNotPersistentFlags}}{{end}}{{ if $explicitlyExposedFlags.HasFlags}}{{flagsUsages $explicitlyExposedFlags}}{{end}}
{{end}}{{ if not $isRootCmd}}Use "{{$rootCmd}} --help" for a list of all commands available in {{$rootCmd}}.
{{end}}{{ if .HasSubCommands }}Use "{{$rootCmd}} <command> --help" for more information about a given command.
{{end}}{{ if and .HasInheritedFlags (not $isRootCmd)}}Use "{{$rootCmd}} options" for a list of global command-line options (applies to all commands).
{{end}}`

	optionsHelpTemplate = ``

	optionsUsageTemplate = `{{ if .HasInheritedFlags}}The following options can be passed to any command:

{{flagsUsages .InheritedFlags}}{{end}}`
)
<|endoftext|>


<commit_before>package ir

import (
	"fmt"

	"github.com/tisp-lang/tisp/src/lib/core"
)

// Switch represents a switch expression.
type Switch struct {
	value       interface{}
	cases       []Case
	defaultCase interface{}
	dict        *core.Thunk
}

// NewSwitch creates a switch expression.
func NewSwitch(v interface{}, cs []Case, d interface{}) Switch {
	if len(cs) == 0 && d == nil {
		panic(fmt.Errorf("A number of cases in switch expressions must be more than 0"))
	}

	return Switch{v, cs, d, compileCasesToDict(cs)}
}

func compileCasesToDict(cs []Case) *core.Thunk {
	ks := make([]core.Value, 0, len(cs))
	vs := make([]*core.Thunk, 0, len(cs))

	for i, c := range cs {
		ks = append(ks, c.pattern.Eval())
		vs = append(vs, core.NewNumber(float64(i)))
	}

	return core.NewDictionary(ks, vs)
}
<commit_msg>Evaluate switch expression dictionary ahead<commit_after>package ir

import (
	"fmt"

	"github.com/tisp-lang/tisp/src/lib/core"
)

// Switch represents a switch expression.
type Switch struct {
	value       interface{}
	cases       []Case
	defaultCase interface{}
	dict        *core.Thunk
}

// NewSwitch creates a switch expression.
func NewSwitch(v interface{}, cs []Case, d interface{}) Switch {
	if len(cs) == 0 && d == nil {
		panic(fmt.Errorf("A number of cases in switch expressions must be more than 0"))
	}

	return Switch{v, cs, d, compileCasesToDict(cs)}
}

func compileCasesToDict(cs []Case) *core.Thunk {
	ks := make([]core.Value, 0, len(cs))
	vs := make([]*core.Thunk, 0, len(cs))

	for i, c := range cs {
		ks = append(ks, c.pattern.Eval())
		vs = append(vs, core.NewNumber(float64(i)))
	}

	return core.Normal(core.NewDictionary(ks, vs).Eval())
}
<|endoftext|>


<commit_before>package v1

import (
	"github.com/fritzpay/paymentd/pkg/paymentd/currency"
	"github.com/fritzpay/paymentd/pkg/service"
	"github.com/gorilla/mux"
	"gopkg.in/inconshreveable/log15.v2"
	"net/http"
)

type CurrencyAdminAPIResponse struct {
	AdminAPIResponse
}

// return a handler brokering get all currencies
func (a *AdminAPI) CurrencyGetRequest() http.Handler {

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		log := a.log.New(log15.Ctx{"method": "Currency Request"})

		// get param
		vars := mux.Vars(r)
		currencyParam := vars["currencycode"]
		if r.Method != "GET" {
			ErrInval.Write(w)
			log.Info("unsupported method " + r.Method)
		}
		log.Info("param: " + currencyParam)

		// get one Currency
		if len(currencyParam) != 3 {
			ErrReadParam.Write(w)
			log.Info("malformed param: " + currencyParam)
			return
		}

		db := a.ctx.PaymentDB(service.ReadOnly)
		c, err := currency.CurrencyByCodeISO4217DB(db, currencyParam)
		if err == currency.ErrCurrencyNotFound {
			ErrNotFound.Write(w)
			log.Info("currency " + currencyParam + " not found")
			return
		} else if err != nil {
			ErrDatabase.Write(w)
			log.Error("database error ", log15.Ctx{"err": err})
			return
		}

		resp := CurrencyAdminAPIResponse{}
		resp.Info = "currency " + c.CodeISO4217 + " found"
		resp.Status = StatusSuccess
		resp.Response = c
		// response write
		resp.Write(w)
		if err != nil {
			log.Error("write error", log15.Ctx{"err": err})
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

	})
}

// return a handler brokering get a currency
func (a *AdminAPI) CurrencyGetAllRequest() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		// get all
		log := a.log.New(log15.Ctx{"method": "Currency Request"})
		db := a.ctx.PaymentDB(service.ReadOnly)
		cl, err := currency.CurrencyAllDB(db)
		if err != nil {
			ErrDatabase.Write(w)
			log.Error("database error", log15.Ctx{"err": err})
			return
		}
		// response write
		resp := CurrencyAdminAPIResponse{}
		resp.Status = StatusSuccess
		resp.Info = "currencies found"
		resp.Response = cl
		resp.Write(w)
		if err != nil {
			log.Error("write error", log15.Ctx{"err": err})
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		return
	})
}
<commit_msg>clean up currency<commit_after>package v1

import (
	"github.com/fritzpay/paymentd/pkg/paymentd/currency"
	"github.com/fritzpay/paymentd/pkg/service"
	"github.com/gorilla/mux"
	"gopkg.in/inconshreveable/log15.v2"
	"net/http"
)

type CurrencyAdminAPIResponse struct {
	AdminAPIResponse
}

// return a handler brokering get all currencies
func (a *AdminAPI) CurrencyGetRequest() http.Handler {

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		log := a.log.New(log15.Ctx{"method": "CurrencyGetRequest"})

		// get param
		vars := mux.Vars(r)
		currencyParam := vars["currencycode"]
		if r.Method != "GET" {
			ErrInval.Write(w)
			log.Info("unsupported method", log15.Ctx{"requestMethod": r.Method})
			return
		}

		// get one Currency
		if len(currencyParam) != 3 {
			ErrReadParam.Write(w)
			log.Info("malformed param", log15.Ctx{"currencyParam": currencyParam})
			return
		}

		log = log.New(log15.Ctx{"currencyParam": currencyParam})

		db := a.ctx.PaymentDB(service.ReadOnly)
		c, err := currency.CurrencyByCodeISO4217DB(db, currencyParam)
		if err == currency.ErrCurrencyNotFound {
			ErrNotFound.Write(w)
			log.Info("currency not found")
			return
		} else if err != nil {
			ErrDatabase.Write(w)
			log.Error("database error", log15.Ctx{"err": err})
			return
		}

		resp := CurrencyAdminAPIResponse{}
		resp.Info = "currency " + c.CodeISO4217 + " found"
		resp.Status = StatusSuccess
		resp.Response = c
		// response write
		resp.Write(w)
		if err != nil {
			log.Error("write error", log15.Ctx{"err": err})
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
	})
}

// return a handler brokering get a currency
func (a *AdminAPI) CurrencyGetAllRequest() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		// get all
		log := a.log.New(log15.Ctx{"method": "Currency Request"})
		db := a.ctx.PaymentDB(service.ReadOnly)
		cl, err := currency.CurrencyAllDB(db)
		if err != nil {
			ErrDatabase.Write(w)
			log.Error("database error", log15.Ctx{"err": err})
			return
		}
		// response write
		resp := CurrencyAdminAPIResponse{}
		resp.Status = StatusSuccess
		resp.Info = "currencies found"
		resp.Response = cl
		resp.Write(w)
		if err != nil {
			log.Error("write error", log15.Ctx{"err": err})
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		return
	})
}
<|endoftext|>


<commit_before>package std

import (
	"time"

	"github.com/tisp-lang/tisp/src/lib/core"
	"github.com/tisp-lang/tisp/src/lib/systemt"
)

const maxConcurrency = 256
const valueChannelCapacity = 1024
const channelCloseDuration = 100 * time.Millisecond

// Rally sorts arguments by time.
var Rally = core.NewLazyFunction(
	core.NewSignature(
		nil, nil, "xs",
		nil, nil, "",
	),
	func(ts ...*core.Thunk) core.Value {
		vs := make(chan core.Value, valueChannelCapacity)

		systemt.Daemonize(func() {
			l := ts[0]
			sem := make(chan bool, maxConcurrency)

			for {
				v := core.PApp(core.Equal, l, core.EmptyList).Eval()
				b, ok := v.(core.BoolType)

				if !ok {
					vs <- core.NotBoolError(v).Eval()
					break
				} else if b {
					time.Sleep(channelCloseDuration)
					vs <- nil
					break
				}

				sem <- true
				go func(t *core.Thunk) {
					vs <- t.Eval()
					<-sem
				}(core.PApp(core.First, l))

				l = core.PApp(core.Rest, l)
			}
		})

		return core.PApp(core.PApp(Y, core.NewLazyFunction(
			core.NewSignature([]string{"me"}, nil, "", nil, nil, ""),
			func(ts ...*core.Thunk) core.Value {
				v := <-vs

				if v == nil {
					return core.EmptyList
				} else if err, ok := v.(core.ErrorType); ok {
					return err
				}

				return core.PApp(core.Prepend, core.Normal(v), core.PApp(ts[0]))
			})))
	})
<commit_msg>Add comment on duration to wait goroutines in rally function<commit_after>package std

import (
	"time"

	"github.com/tisp-lang/tisp/src/lib/core"
	"github.com/tisp-lang/tisp/src/lib/systemt"
)

const maxConcurrency = 256
const valueChannelCapacity = 1024
const channelCloseDuration = 100 * time.Millisecond

// Rally sorts arguments by time.
var Rally = core.NewLazyFunction(
	core.NewSignature(
		nil, nil, "xs",
		nil, nil, "",
	),
	func(ts ...*core.Thunk) core.Value {
		vs := make(chan core.Value, valueChannelCapacity)

		systemt.Daemonize(func() {
			l := ts[0]
			sem := make(chan bool, maxConcurrency)

			for {
				v := core.PApp(core.Equal, l, core.EmptyList).Eval()
				b, ok := v.(core.BoolType)

				if !ok {
					vs <- core.NotBoolError(v).Eval()
					break
				} else if b {
					// HACK: Wait for other goroutines to put elements in a value channel
					// for a while. This is only for unit test.
					time.Sleep(channelCloseDuration)
					vs <- nil
					break
				}

				sem <- true
				go func(t *core.Thunk) {
					vs <- t.Eval()
					<-sem
				}(core.PApp(core.First, l))

				l = core.PApp(core.Rest, l)
			}
		})

		return core.PApp(core.PApp(Y, core.NewLazyFunction(
			core.NewSignature([]string{"me"}, nil, "", nil, nil, ""),
			func(ts ...*core.Thunk) core.Value {
				v := <-vs

				if v == nil {
					return core.EmptyList
				} else if err, ok := v.(core.ErrorType); ok {
					return err
				}

				return core.PApp(core.Prepend, core.Normal(v), core.PApp(ts[0]))
			})))
	})
<|endoftext|>


<commit_before>package sqlstore

import (
	"context"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
)

func init() {
	bus.AddHandler("sql", GetSystemStats)
	bus.AddHandler("sql", GetDataSourceStats)
	bus.AddHandler("sql", GetDataSourceAccessStats)
	bus.AddHandler("sql", GetAdminStats)
	bus.AddHandlerCtx("sql", GetAlertNotifiersUsageStats)
	bus.AddHandlerCtx("sql", GetSystemUserCountStats)
}

var activeUserTimeLimit = time.Hour * 24 * 30

func GetAlertNotifiersUsageStats(ctx context.Context, query *m.GetAlertNotifierUsageStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type FROM alert_notification GROUP BY type`
	query.Result = make([]*m.NotifierUsageStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetDataSourceStats(query *m.GetDataSourceStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type FROM data_source GROUP BY type`
	query.Result = make([]*m.DataSourceStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetDataSourceAccessStats(query *m.GetDataSourceAccessStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type, access FROM data_source GROUP BY type, access`
	query.Result = make([]*m.DataSourceAccessStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetSystemStats(query *m.GetSystemStatsQuery) error {
	sb := &SqlBuilder{}
	sb.Write("SELECT ")
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("user") + `) AS users,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("org") + `) AS orgs,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("dashboard") + `) AS dashboards,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("data_source") + `) AS datasources,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("star") + `) AS stars,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("playlist") + `) AS playlists,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("alert") + `) AS alerts,`)

	activeUserDeadlineDate := time.Now().Add(-activeUserTimeLimit)
	sb.Write(`(SELECT COUNT(*) FROM `+dialect.Quote("user")+` where last_seen_at > ?) AS active_users,`, activeUserDeadlineDate)

	sb.Write(`(SELECT COUNT(id) FROM `+dialect.Quote("dashboard")+` where is_folder = ?) AS folders,`, dialect.BooleanStr(true))

	sb.Write(`(
		SELECT COUNT(acl.id)
		FROM `+dialect.Quote("dashboard_acl")+` as acl
			inner join `+dialect.Quote("dashboard")+` as d
			on d.id = acl.dashboard_id
		WHERE d.is_folder = ?
	) AS dashboard_permissions,`, dialect.BooleanStr(false))

	sb.Write(`(
		SELECT COUNT(acl.id)
		FROM `+dialect.Quote("dashboard_acl")+` as acl
			inner join `+dialect.Quote("dashboard")+` as d
			on d.id = acl.dashboard_id
		WHERE d.is_folder = ?
	) AS folder_permissions,`, dialect.BooleanStr(true))

	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("dashboard_provisioning") + `) AS provisioned_dashboards,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("dashboard_snapshot") + `) AS snapshots,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("team") + `) AS teams,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("user_auth_token") + `) AS auth_tokens,`)

	sb.Write(roleCounterSQL("Viewer", "viewers")+`,`, activeUserDeadlineDate)
	sb.Write(roleCounterSQL("Editor", "editors")+`,`, activeUserDeadlineDate)
	sb.Write(roleCounterSQL("Admin", "admins")+``, activeUserDeadlineDate)

	var stats m.SystemStats
	_, err := x.SQL(sb.GetSqlString(), sb.params...).Get(&stats)
	if err != nil {
		return err
	}

	query.Result = &stats

	return err
}

func roleCounterSQL(role, alias string) string {
	return `
		(
			SELECT COUNT(*)
			FROM ` + dialect.Quote("user") + ` as u, org_user
			WHERE ( org_user.user_id=u.id AND org_user.role='` + role + `' )
		) as ` + alias + `,
		(
			SELECT COUNT(*)
			FROM ` + dialect.Quote("user") + ` as u, org_user
			WHERE u.last_seen_at>? AND ( org_user.user_id=u.id AND org_user.role='` + role + `' )
		) as active_` + alias
}

func GetAdminStats(query *m.GetAdminStatsQuery) error {
	activeEndDate := time.Now().Add(-activeUserTimeLimit)

	var rawSql = `SELECT
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("org") + `
		  ) AS orgs,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("dashboard") + `
		) AS dashboards,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("dashboard_snapshot") + `
		  ) AS snapshots,
		  (
		SELECT COUNT( DISTINCT ( ` + dialect.Quote("term") + ` ))
		FROM ` + dialect.Quote("dashboard_tag") + `
		  ) AS tags,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("data_source") + `
		  ) AS datasources,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("playlist") + `
		  ) AS playlists,
		  (
		SELECT COUNT(*) FROM ` + dialect.Quote("star") + `
		  ) AS stars,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("alert") + `
		) AS alerts,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user") + `
		) AS users,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user") + ` where last_seen_at > ?
		) as active_users,
		` + roleCounterSQL("Admin", "admins") + `,
		` + roleCounterSQL("Editor", "editors") + `,
		` + roleCounterSQL("Viewer", "viewers") + `,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user_auth_token") + ` where rotated_at > ?
		) as active_sessions
	  `

	var stats m.AdminStats
	_, err := x.SQL(rawSql, activeEndDate, activeEndDate, activeEndDate, activeEndDate, activeEndDate.Unix()).Get(&stats)
	if err != nil {
		return err
	}

	query.Result = &stats
	return err
}

func GetSystemUserCountStats(ctx context.Context, query *m.GetSystemUserCountStatsQuery) error {
	return withDbSession(ctx, func(sess *DBSession) error {

		var rawSql = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
		var stats m.SystemUserCountStats
		_, err := sess.SQL(rawSql).Get(&stats)
		if err != nil {
			return err
		}

		query.Result = &stats

		return err
	})
}
<commit_msg>SQLStore: Rewrite system statistics query to count users once (#20711)<commit_after>package sqlstore

import (
	"context"
	"time"

	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
)

func init() {
	bus.AddHandler("sql", GetSystemStats)
	bus.AddHandler("sql", GetDataSourceStats)
	bus.AddHandler("sql", GetDataSourceAccessStats)
	bus.AddHandler("sql", GetAdminStats)
	bus.AddHandlerCtx("sql", GetAlertNotifiersUsageStats)
	bus.AddHandlerCtx("sql", GetSystemUserCountStats)
}

var activeUserTimeLimit = time.Hour * 24 * 30

func GetAlertNotifiersUsageStats(ctx context.Context, query *m.GetAlertNotifierUsageStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type FROM alert_notification GROUP BY type`
	query.Result = make([]*m.NotifierUsageStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetDataSourceStats(query *m.GetDataSourceStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type FROM data_source GROUP BY type`
	query.Result = make([]*m.DataSourceStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetDataSourceAccessStats(query *m.GetDataSourceAccessStatsQuery) error {
	var rawSql = `SELECT COUNT(*) as count, type, access FROM data_source GROUP BY type, access`
	query.Result = make([]*m.DataSourceAccessStats, 0)
	err := x.SQL(rawSql).Find(&query.Result)
	return err
}

func GetSystemStats(query *m.GetSystemStatsQuery) error {
	sb := &SqlBuilder{}
	sb.Write("SELECT ")
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("user") + `) AS users,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("org") + `) AS orgs,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("dashboard") + `) AS dashboards,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("data_source") + `) AS datasources,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("star") + `) AS stars,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("playlist") + `) AS playlists,`)
	sb.Write(`(SELECT COUNT(*) FROM ` + dialect.Quote("alert") + `) AS alerts,`)

	activeUserDeadlineDate := time.Now().Add(-activeUserTimeLimit)
	sb.Write(`(SELECT COUNT(*) FROM `+dialect.Quote("user")+` where last_seen_at > ?) AS active_users,`, activeUserDeadlineDate)

	sb.Write(`(SELECT COUNT(id) FROM `+dialect.Quote("dashboard")+` where is_folder = ?) AS folders,`, dialect.BooleanStr(true))

	sb.Write(`(
		SELECT COUNT(acl.id)
		FROM `+dialect.Quote("dashboard_acl")+` as acl
			inner join `+dialect.Quote("dashboard")+` as d
			on d.id = acl.dashboard_id
		WHERE d.is_folder = ?
	) AS dashboard_permissions,`, dialect.BooleanStr(false))

	sb.Write(`(
		SELECT COUNT(acl.id)
		FROM `+dialect.Quote("dashboard_acl")+` as acl
			inner join `+dialect.Quote("dashboard")+` as d
			on d.id = acl.dashboard_id
		WHERE d.is_folder = ?
	) AS folder_permissions,`, dialect.BooleanStr(true))

	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("dashboard_provisioning") + `) AS provisioned_dashboards,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("dashboard_snapshot") + `) AS snapshots,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("team") + `) AS teams,`)
	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("user_auth_token") + `) AS auth_tokens,`)

	sb.Write(roleCounterSQL("Viewer", "viewers")+`,`, activeUserDeadlineDate)
	sb.Write(roleCounterSQL("Editor", "editors")+`,`, activeUserDeadlineDate)
	sb.Write(roleCounterSQL("Admin", "admins")+``, activeUserDeadlineDate)

	var stats m.SystemStats
	_, err := x.SQL(sb.GetSqlString(), sb.params...).Get(&stats)
	if err != nil {
		return err
	}

	query.Result = &stats

	return err
}

func roleCounterSQL(role, alias string) string {
	return `
		(
			SELECT COUNT(DISTINCT u.id)
			FROM ` + dialect.Quote("user") + ` as u, org_user
			WHERE ( org_user.user_id=u.id AND org_user.role='` + role + `' )
		) as ` + alias + `,
		(
			SELECT COUNT(DISTINCT u.id)
			FROM ` + dialect.Quote("user") + ` as u, org_user
			WHERE u.last_seen_at>? AND ( org_user.user_id=u.id AND org_user.role='` + role + `' )
		) as active_` + alias
}

func GetAdminStats(query *m.GetAdminStatsQuery) error {
	activeEndDate := time.Now().Add(-activeUserTimeLimit)

	var rawSql = `SELECT
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("org") + `
		  ) AS orgs,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("dashboard") + `
		) AS dashboards,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("dashboard_snapshot") + `
		  ) AS snapshots,
		  (
		SELECT COUNT( DISTINCT ( ` + dialect.Quote("term") + ` ))
		FROM ` + dialect.Quote("dashboard_tag") + `
		  ) AS tags,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("data_source") + `
		  ) AS datasources,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("playlist") + `
		  ) AS playlists,
		  (
		SELECT COUNT(*) FROM ` + dialect.Quote("star") + `
		  ) AS stars,
		  (
		SELECT COUNT(*)
		FROM ` + dialect.Quote("alert") + `
		) AS alerts,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user") + `
		) AS users,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user") + ` where last_seen_at > ?
		) as active_users,
		` + roleCounterSQL("Admin", "admins") + `,
		` + roleCounterSQL("Editor", "editors") + `,
		` + roleCounterSQL("Viewer", "viewers") + `,
		(
		SELECT COUNT(*)
		FROM ` + dialect.Quote("user_auth_token") + ` where rotated_at > ?
		) as active_sessions
	  `

	var stats m.AdminStats
	_, err := x.SQL(rawSql, activeEndDate, activeEndDate, activeEndDate, activeEndDate, activeEndDate.Unix()).Get(&stats)
	if err != nil {
		return err
	}

	query.Result = &stats
	return err
}

func GetSystemUserCountStats(ctx context.Context, query *m.GetSystemUserCountStatsQuery) error {
	return withDbSession(ctx, func(sess *DBSession) error {

		var rawSql = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
		var stats m.SystemUserCountStats
		_, err := sess.SQL(rawSql).Get(&stats)
		if err != nil {
			return err
		}

		query.Result = &stats

		return err
	})
}
<|endoftext|>