<commit_before>//Copyright (c) 2013 The Go-IMAP Authors. All rights reserved.
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions are
//met:
//* Redistributions of source code must retain the above copyright
//notice, this list of conditions and the following disclaimer.
//* Redistributions in binary form must reproduce the above copyright
//notice, this list of conditions and the following disclaimer in the
//documentation and/or other materials provided with the
//distribution.
//* Neither the name of the go-imap project nor the names of its
//contributors may be used to endorse or promote products derived
//from this software without specific prior written permission.
//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
//"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
//LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
//A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
//OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
//SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
//LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
//DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
//THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
//(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
//OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package imapclient

import (
	"github.com/mxk/go-imap/imap"
)

//Login helper functions derived from go-imap demo1.go

func loginIMAP(c *imap.Client, user, pass string) (cmd *imap.Command, err error) {
	//defer c.SetLogMask(sensitiveLogin(c, "LOGIN"))
	return c.Login(user, pass)
}

func loginCramIMAP(c *imap.Client, user, pass string) (cmd *imap.Command, err error) {
	//defer c.SetLogMask(sensitiveLogin(c, "AUTHENTICATE"))

	return c.Auth(CRAMMD5Auth(user, pass))
}

func sensitiveLogin(c *imap.Client, action string) imap.LogMask {
	mask := c.SetLogMask(imap.LogConn)
	hide := imap.LogCmd | imap.LogRaw
	if mask&hide != 0 {
		c.Logln(imap.LogConn, "Raw logging disabled during", action)
	}
	c.SetLogMask(mask &^ hide)
	return mask
}
<commit_msg>Do not log IMAP protocol authentication<commit_after>//Copyright (c) 2013 The Go-IMAP Authors. All rights reserved.
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions are
//met:
//* Redistributions of source code must retain the above copyright
//notice, this list of conditions and the following disclaimer.
//* Redistributions in binary form must reproduce the above copyright
//notice, this list of conditions and the following disclaimer in the
//documentation and/or other materials provided with the
//distribution.
//* Neither the name of the go-imap project nor the names of its
//contributors may be used to endorse or promote products derived
//from this software without specific prior written permission.
//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
//"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
//LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
//A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
//OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
//SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
//LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
//DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
//THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
//(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
//OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package imapclient

import (
	"github.com/mxk/go-imap/imap"
)

//Login helper functions derived from go-imap demo1.go

func loginIMAP(c *imap.Client, user, pass string) (cmd *imap.Command, err error) {
	defer c.SetLogMask(sensitiveLogin(c, "LOGIN"))
	return c.Login(user, pass)
}

func loginCramIMAP(c *imap.Client, user, pass string) (cmd *imap.Command, err error) {
	defer c.SetLogMask(sensitiveLogin(c, "AUTHENTICATE"))

	return c.Auth(CRAMMD5Auth(user, pass))
}

func sensitiveLogin(c *imap.Client, action string) imap.LogMask {
	mask := c.SetLogMask(imap.LogConn)
	hide := imap.LogCmd | imap.LogRaw
	if mask&hide != 0 {
		c.Logln(imap.LogConn, "Raw logging disabled during", action)
	}
	c.SetLogMask(mask &^ hide)
	return mask
}
<|endoftext|>


<commit_before>// Copyright 2015 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.

// Helper code for daemonizing gcsfuse, synchronizing on successful mount.
//
// The details of this package are subject to change.
package daemon

import (
	"encoding/gob"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
)

// The name of an environment variable used to communicate a file descriptor
// set up by Mount to the gcsfuse subprocess. Gob encoding is used to
// communicate back to Mount.
const envVar = "MOUNT_STATUS_FD"

// A message containing logging output for the process of mounting the file
// system.
type logMsg struct {
	Msg []byte
}

// A message indicating the outcome of the process of mounting the file system.
// The receiver ignores further messages.
type outcomeMsg struct {
	Succesful bool

	// Meaningful only if !Succesful.
	ErrorMsg string
}

// The file provded to this process via the environment variable, or nil if
// none.
var gFile *os.File

// A gob encoder that writes into gFile, or nil.
var gGobEncoder *gob.Encoder

func init() {
	// Is the environment variable set?
	fdStr, ok := os.LookupEnv(envVar)
	if !ok {
		return
	}

	// Parse the file descriptor.
	fd, err := strconv.ParseUint(fdStr, 10, 32)
	if err != nil {
		log.Fatalf("Couldn't parse %s value %q: %v", envVar, fdStr, err)
	}

	// Set up the file and the encoder that wraps it.
	gFile = os.NewFile(uintptr(fd), envVar)
	gGobEncoder = gob.NewEncoder(gFile)
}

// For use by gcsfuse: signal that mounting was successful (allowing the caller
// of the process to return in success) or that there was a failure to mount
// the file system (allowing the caller of the process to display an
// appropriate error message).
//
// Do nothing if the process wasn't invoked with Mount.
func SignalOutcome(outcome error) (err error) {
	// Is there anything to do?
	if gGobEncoder == nil {
		return
	}

	// Write out the outcome.
	msg := &outcomeMsg{
		Succesful: outcome == nil,
	}

	if !msg.Succesful {
		msg.ErrorMsg = outcome.Error()
	}

	err = gGobEncoder.Encode(msg)
	return
}

// An io.Writer that sends logMsg messages over a gob encoder.
type logMsgWriter struct {
	encoder *gob.Encoder
}

func (w *logMsgWriter) Write(p []byte) (n int, err error) {
	msg := &logMsg{
		Msg: p,
	}

	err = w.encoder.Encode(msg)
	if err != nil {
		return
	}

	n = len(p)
	return
}

// For use by gcsfuse: return a writer that should be used for logging status
// messages while in the process of mounting.
//
// The returned writer must not be written to after calling SignalOutcome.
func StatusWriter() (w io.Writer) {
	if gGobEncoder == nil {
		w = ioutil.Discard
		return
	}

	w = &logMsgWriter{gGobEncoder}
	return
}

// Invoke gcsfuse with the supplied arguments, waiting until it successfully
// mounts or reports that is has failed. Write status updates while mounting
// into the supplied writer (which may be nil for silence). Return nil only if
// it mounts successfully.
func Mount(
	gcsfusePath string,
	args []string,
	status io.Writer) (err error) {
	err = errors.New("TODO")
	return
}
<commit_msg>daemon.Mount<commit_after>// Copyright 2015 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.

// Helper code for daemonizing gcsfuse, synchronizing on successful mount.
//
// The details of this package are subject to change.
package daemon

import (
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
)

// The name of an environment variable used to communicate a file descriptor
// set up by Mount to the gcsfuse subprocess. Gob encoding is used to
// communicate back to Mount.
const envVar = "MOUNT_STATUS_FD"

// A message containing logging output for the process of mounting the file
// system.
type logMsg struct {
	Msg []byte
}

// A message indicating the outcome of the process of mounting the file system.
// The receiver ignores further messages.
type outcomeMsg struct {
	Succesful bool

	// Meaningful only if !Succesful.
	ErrorMsg string
}

// The file provded to this process via the environment variable, or nil if
// none.
var gFile *os.File

// A gob encoder that writes into gFile, or nil.
var gGobEncoder *gob.Encoder

func init() {
	// Is the environment variable set?
	fdStr, ok := os.LookupEnv(envVar)
	if !ok {
		return
	}

	// Parse the file descriptor.
	fd, err := strconv.ParseUint(fdStr, 10, 32)
	if err != nil {
		log.Fatalf("Couldn't parse %s value %q: %v", envVar, fdStr, err)
	}

	// Set up the file and the encoder that wraps it.
	gFile = os.NewFile(uintptr(fd), envVar)
	gGobEncoder = gob.NewEncoder(gFile)
}

// For use by gcsfuse: signal that mounting was successful (allowing the caller
// of the process to return in success) or that there was a failure to mount
// the file system (allowing the caller of the process to display an
// appropriate error message).
//
// Do nothing if the process wasn't invoked with Mount.
func SignalOutcome(outcome error) (err error) {
	// Is there anything to do?
	if gGobEncoder == nil {
		return
	}

	// Write out the outcome.
	msg := &outcomeMsg{
		Succesful: outcome == nil,
	}

	if !msg.Succesful {
		msg.ErrorMsg = outcome.Error()
	}

	err = gGobEncoder.Encode(msg)
	return
}

// An io.Writer that sends logMsg messages over a gob encoder.
type logMsgWriter struct {
	encoder *gob.Encoder
}

func (w *logMsgWriter) Write(p []byte) (n int, err error) {
	msg := &logMsg{
		Msg: p,
	}

	err = w.encoder.Encode(msg)
	if err != nil {
		return
	}

	n = len(p)
	return
}

// For use by gcsfuse: return a writer that should be used for logging status
// messages while in the process of mounting.
//
// The returned writer must not be written to after calling SignalOutcome.
func StatusWriter() (w io.Writer) {
	if gGobEncoder == nil {
		w = ioutil.Discard
		return
	}

	w = &logMsgWriter{gGobEncoder}
	return
}

// Invoke gcsfuse with the supplied arguments, waiting until it successfully
// mounts or reports that is has failed. Write status updates while mounting
// into the supplied writer (which may be nil for silence). Return nil only if
// it mounts successfully.
func Mount(
	gcsfusePath string,
	args []string,
	status io.Writer) (err error) {
	if status == nil {
		status = ioutil.Discard
	}

	// Set up the pipe that we will hand to the gcsfuse subprocess.
	pipeR, pipeW, err := os.Pipe()
	if err != nil {
		err = fmt.Errorf("Pipe: %v", err)
		return
	}

	// Attempt to start gcsfuse. If we encounter an error in so doing, write it
	// to the channel.
	startGcsfuseErr := make(chan error, 1)
	go func() {
		defer pipeW.Close()
		startGcsfuseErr <- startGcsfuse(args, pipeW)
	}()

	// Read communication from gcsfuse from the pipe, writing nil into the
	// channel only if the mount succeeds.
	readFromGcsfuseOutcome := make(chan error, 1)
	go func() {
		defer pipeR.Close()
		readFromGcsfuseOutcome <- readFromGcsfuse(pipeR, status)
	}()

	// Wait for a result from one of the above.
	select {
	case err = <-startGcsfuseErr:
		err = fmt.Errorf("startGcsfuse: %v", err)
		return

	case err = <-readFromGcsfuseOutcome:
		if err == nil {
			// All is good.
			return
		}

		err = fmt.Errorf("readFromGcsfuse: %v", err)
		return
	}
}

// Start gcsfuse, handing it the supplied pipe for communication. Do not wait
// for it to return.
func startGcsfuse(
	args []string,
	pipeW *os.File) (err error) {
	err = errors.New("TODO")
	return
}

// Process communication from a gcsfuse subprocess. Write log messages to the
// supplied writer (which must be non-nil). Return nil only if the output of
// mounting is success.
func readFromGcsfuse(
	r io.Reader,
	status io.Writer) (err error) {
	err = errors.New("TODO")
	return
}
<|endoftext|>


<commit_before>// Copyright 2018 The Go 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 web

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"syscall"
	"testing"
	"testing/fstest"

	"github.com/google/go-cmp/cmp"
)

func testServeBody(t *testing.T, p *Site, path, body string) {
	t.Helper()
	r := &http.Request{URL: &url.URL{Path: path}}
	rw := httptest.NewRecorder()
	p.ServeHTTP(rw, r)
	if rw.Code != 200 || !strings.Contains(rw.Body.String(), body) {
		t.Fatalf("GET %s: expected 200 w/ %q: got %d w/ body:\n%s",
			path, body, rw.Code, rw.Body)
	}
}

func TestRedirectAndMetadata(t *testing.T) {
	fsys := fstest.MapFS{
		"site.tmpl":           {Data: []byte(`{{.Content}}`)},
		"doc/x/index.html":    {Data: []byte("Hello, x.")},
		"lib/godoc/site.html": {Data: []byte(`{{.Data}}`)},
	}
	site := NewSite(fsys)

	// Test that redirect is sent back correctly.
	// Used to panic. See golang.org/issue/40665.
	dir := "/doc/x/"

	r := &http.Request{URL: &url.URL{Path: dir + "index.html"}}
	rw := httptest.NewRecorder()
	site.ServeHTTP(rw, r)
	loc := rw.Result().Header.Get("Location")
	if rw.Code != 301 || loc != dir {
		t.Errorf("GET %s: expected 301 -> %q, got %d -> %q", r.URL.Path, dir, rw.Code, loc)
	}

	testServeBody(t, site, dir, "Hello, x")
}

func TestMarkdown(t *testing.T) {
	site := NewSite(fstest.MapFS{
		"site.tmpl":           {Data: []byte(`{{.Content}}`)},
		"doc/test.md":         {Data: []byte("**bold**")},
		"doc/test2.md":        {Data: []byte(`{{"*template*"}}`)},
		"lib/godoc/site.html": {Data: []byte(`{{.Data}}`)},
	})

	testServeBody(t, site, "/doc/test", "bold")
	testServeBody(t, site, "/doc/test2", "template")
}

func TestTypeScript(t *testing.T) {
	exampleOut, err := os.ReadFile("testdata/example.js")
	if err != nil {
		t.Fatal(err)
	}
	tests := []struct {
		name            string
		path            string
		wantCode        int
		wantBody        string
		wantCacheHeader bool
	}{
		{
			name:     "example code",
			path:     "/example.ts",
			wantCode: 200,
			wantBody: string(exampleOut),
		},
		{
			name:            "example code cached",
			path:            "/example.ts",
			wantCode:        200,
			wantBody:        string(exampleOut),
			wantCacheHeader: true,
		},
		{
			name:     "file not found",
			path:     "/notfound.ts",
			wantCode: 500,
			wantBody: fmt.Sprintf("\n\nopen testdata/notfound.ts: %s\n", syscall.ENOENT),
		},
		{
			name:     "syntax error",
			path:     "/error.ts",
			wantCode: 500,
			wantBody: "\n\nExpected identifier but found "function"\n\n",
		},
	}
	fsys := os.DirFS("testdata")
	site := NewSite(fsys)
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, tt.path, nil)
			got := httptest.NewRecorder()
			site.serveTypeScript(got, req)
			gotHeader := got.Header().Get(cacheHeader) == "true"
			if got.Code != tt.wantCode {
				t.Errorf("got status %d but wanted %d", got.Code, http.StatusOK)
			}
			if (tt.wantCacheHeader && !gotHeader) || (!tt.wantCacheHeader && gotHeader) {
				t.Errorf("got cache hit %v but wanted %v", gotHeader, tt.wantCacheHeader)
			}
			if diff := cmp.Diff(tt.wantBody, got.Body.String()); diff != "" {
				t.Errorf("ServeHTTP() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}
<commit_msg>internal/web: adjust test to not require exact error message<commit_after>// Copyright 2018 The Go 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 web

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"syscall"
	"testing"
	"testing/fstest"

	"github.com/google/go-cmp/cmp"
)

func testServeBody(t *testing.T, p *Site, path, body string) {
	t.Helper()
	r := &http.Request{URL: &url.URL{Path: path}}
	rw := httptest.NewRecorder()
	p.ServeHTTP(rw, r)
	if rw.Code != 200 || !strings.Contains(rw.Body.String(), body) {
		t.Fatalf("GET %s: expected 200 w/ %q: got %d w/ body:\n%s",
			path, body, rw.Code, rw.Body)
	}
}

func TestRedirectAndMetadata(t *testing.T) {
	fsys := fstest.MapFS{
		"site.tmpl":           {Data: []byte(`{{.Content}}`)},
		"doc/x/index.html":    {Data: []byte("Hello, x.")},
		"lib/godoc/site.html": {Data: []byte(`{{.Data}}`)},
	}
	site := NewSite(fsys)

	// Test that redirect is sent back correctly.
	// Used to panic. See golang.org/issue/40665.
	dir := "/doc/x/"

	r := &http.Request{URL: &url.URL{Path: dir + "index.html"}}
	rw := httptest.NewRecorder()
	site.ServeHTTP(rw, r)
	loc := rw.Result().Header.Get("Location")
	if rw.Code != 301 || loc != dir {
		t.Errorf("GET %s: expected 301 -> %q, got %d -> %q", r.URL.Path, dir, rw.Code, loc)
	}

	testServeBody(t, site, dir, "Hello, x")
}

func TestMarkdown(t *testing.T) {
	site := NewSite(fstest.MapFS{
		"site.tmpl":           {Data: []byte(`{{.Content}}`)},
		"doc/test.md":         {Data: []byte("**bold**")},
		"doc/test2.md":        {Data: []byte(`{{"*template*"}}`)},
		"lib/godoc/site.html": {Data: []byte(`{{.Data}}`)},
	})

	testServeBody(t, site, "/doc/test", "bold")
	testServeBody(t, site, "/doc/test2", "template")
}

func TestTypeScript(t *testing.T) {
	exampleOut, err := os.ReadFile("testdata/example.js")
	if err != nil {
		t.Fatal(err)
	}
	tests := []struct {
		name            string
		path            string
		wantCode        int
		wantBody        string
		wantCacheHeader bool
	}{
		{
			name:     "example code",
			path:     "/example.ts",
			wantCode: 200,
			wantBody: string(exampleOut),
		},
		{
			name:            "example code cached",
			path:            "/example.ts",
			wantCode:        200,
			wantBody:        string(exampleOut),
			wantCacheHeader: true,
		},
		{
			name:     "file not found",
			path:     "/notfound.ts",
			wantCode: 500,
			wantBody: syscall.ENOENT.Error(),
		},
		{
			name:     "syntax error",
			path:     "/error.ts",
			wantCode: 500,
			wantBody: "\n\nExpected identifier but found "function"\n\n",
		},
	}
	fsys := os.DirFS("testdata")
	site := NewSite(fsys)
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, tt.path, nil)
			got := httptest.NewRecorder()
			site.serveTypeScript(got, req)
			gotHeader := got.Header().Get(cacheHeader) == "true"
			if got.Code != tt.wantCode {
				t.Errorf("got status %d but wanted %d", got.Code, http.StatusOK)
			}
			if (tt.wantCacheHeader && !gotHeader) || (!tt.wantCacheHeader && gotHeader) {
				t.Errorf("got cache hit %v but wanted %v", gotHeader, tt.wantCacheHeader)
			}
			if !strings.Contains(got.Body.String(), tt.wantBody) {
				t.Errorf("ServeHTTP() mismatch (-want +got):\n%s", cmp.Diff(tt.wantBody, got.Body.String()))
			}
		})
	}
}
<|endoftext|>


<commit_before>package integration

import (
	"fmt"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/requests"
	"github.com/aliyun/alibaba-cloud-sdk-go/services/ecs"
	"github.com/stretchr/testify/assert"
	"strconv"
	"strings"
	"testing"
	"time"
)

const (
	EcsInstanceDefaultTimeout = 120
	EcsDefaultWaitForInterval = 20

	EcsInstanceStatusRunning = "Running"
	EcsInstanceStatusStopped = "Stopped"
	EcsInstanceStatusDeleted = "Deleted"
)

// create -> start -> stop -> delete
func TestEcsInstance(t *testing.T) {

	// init client
	config := getConfigFromEnv()
	ecsClient, err := ecs.NewClientWithAccessKey("cn-hangzhou", config.AccessKeyId, config.AccessKeySecret)
	assertErrorNil(t, err, "Failed to init client")
	fmt.Printf("Init client success\n")

	// get demo instance attributes
	param := getDemoEcsInstanceAttributes(t, ecsClient)

	// create
	instanceId := createEcsInstance(t, ecsClient, param)

	// defer wait for deleted
	defer waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusDeleted, 120)

	// defer delete
	defer deleteEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusStopped, 60)

	// start
	startEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusRunning, 240)

	// stop
	stopEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusStopped, 600)

	//delete all test instance
	deleteAllTestEcsInstance(t, ecsClient)
}

func getDemoEcsInstanceAttributes(t *testing.T, client *ecs.Client) *ecs.DescribeInstanceAttributeResponse {
	fmt.Print("trying to get demo ecs instance...")
	request := ecs.CreateDescribeInstanceAttributeRequest()
	request.InstanceId = getEcsDemoInstanceId()
	response, err := client.DescribeInstanceAttribute(request)
	assertErrorNil(t, err, "Failed to get demo ecs instance attributes")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
	return response
}

func createEcsInstance(t *testing.T, client *ecs.Client, param *ecs.DescribeInstanceAttributeResponse) (instanceId string) {
	fmt.Print("creating ecs instance...")
	request := ecs.CreateCreateInstanceRequest()
	request.ImageId = param.ImageId
	request.InstanceName = InstanceNamePrefix + strconv.FormatInt(time.Now().Unix(), 10)
	request.SecurityGroupId = param.SecurityGroupIds.SecurityGroupId[0]
	request.InstanceType = "ecs.t1.small"
	response, err := client.CreateInstance(request)
	assertErrorNil(t, err, "Failed to create ecs instance")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	instanceId = response.InstanceId
	fmt.Printf("success(%d)! instanceId = %s\n", response.GetHttpStatus(), instanceId)
	return
}

func startEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("starting ecs instance(%s)...", instanceId)
	request := ecs.CreateStartInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.StartInstance(request)
	assertErrorNil(t, err, "Failed to start ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func stopEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("stopping ecs instance(%s)...", instanceId)
	request := ecs.CreateStopInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.StopInstance(request)
	assertErrorNil(t, err, "Failed to stop ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func deleteEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("deleting ecs instance(%s)...", instanceId)
	request := ecs.CreateDeleteInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.DeleteInstance(request)
	assertErrorNil(t, err, "Failed to delete ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func deleteAllTestEcsInstance(t *testing.T, client *ecs.Client) {
	fmt.Print("list all ecs instances...")
	request := ecs.CreateDescribeInstancesRequest()
	request.PageSize = requests.NewInteger(10)
	request.PageNumber = requests.NewInteger(1)
	response, err := client.DescribeInstances(request)
	assertErrorNil(t, err, "Failed to list all ecs instances ")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Printf("success! TotalCount = %s\n", strconv.Itoa(response.TotalCount))

	for _, instanceInfo := range response.Instances.Instance {
		if strings.HasPrefix(instanceInfo.InstanceName, InstanceNamePrefix) {
			createTime, err := strconv.ParseInt(instanceInfo.InstanceName[26:len(instanceInfo.InstanceName)], 10, 64)
			assertErrorNil(t, err, "Parse instance create time failed: "+instanceInfo.InstanceName)
			if (time.Now().Unix() - createTime) < (60 * 60) {
				fmt.Printf("found undeleted ecs instance(%s) but created in 60 minutes, try to delete next time\n", instanceInfo.InstanceName)
				return
			} else {
				fmt.Printf("found undeleted ecs instance(%s), status=%s, try to delete it.\n",
					instanceInfo.Status, instanceInfo.InstanceId)
				if instanceInfo.Status == EcsInstanceStatusRunning {
					// stop
					stopEcsInstance(t, client, instanceInfo.InstanceId)
				} else if instanceInfo.Status == EcsInstanceStatusStopped {
					// delete
					deleteEcsInstance(t, client, instanceInfo.InstanceId)
					// wait
					waitForEcsInstance(t, client, instanceInfo.InstanceId, EcsInstanceStatusDeleted, 600)
				}
			}
		}
	}
}

func waitForEcsInstance(t *testing.T, client *ecs.Client, instanceId string, targetStatus string, timeout int) {
	if timeout <= 0 {
		timeout = EcsInstanceDefaultTimeout
	}
	for {
		request := ecs.CreateDescribeInstanceAttributeRequest()
		request.InstanceId = instanceId
		response, err := client.DescribeInstanceAttribute(request)

		if targetStatus == EcsInstanceStatusDeleted {
			if response.GetHttpStatus() == 404 || response.Status == EcsInstanceStatusDeleted {
				fmt.Printf("delete ecs instance(%s) success\n", instanceId)
				break
			} else {
				assertErrorNil(t, err, "Failed to describe ecs instance \n")
			}
		} else {
			assertErrorNil(t, err, "Failed to describe ecs instance \n")
			if response.Status == targetStatus {
				fmt.Printf("ecs instance(%s) status changed to %s, wait a moment\n", instanceId, targetStatus)
				time.Sleep(EcsDefaultWaitForInterval * time.Second)
				break
			} else {
				fmt.Printf("ecs instance(%s) status is %s, wait for changing to %s\n", instanceId, response.Status, targetStatus)
			}
		}

		timeout = timeout - EcsDefaultWaitForInterval
		if timeout <= 0 {
			t.Errorf(fmt.Sprintf("wait for ecs instance(%s) status to %s timeout(%d)\n", instanceId, targetStatus, timeout))
		}
		time.Sleep(EcsDefaultWaitForInterval * time.Second)
	}
}
<commit_msg>fix a bug in delete all test instance<commit_after>package integration

import (
	"fmt"
	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/requests"
	"github.com/aliyun/alibaba-cloud-sdk-go/services/ecs"
	"github.com/stretchr/testify/assert"
	"strconv"
	"strings"
	"testing"
	"time"
)

const (
	EcsInstanceDefaultTimeout = 120
	EcsDefaultWaitForInterval = 20

	EcsInstanceStatusRunning = "Running"
	EcsInstanceStatusStopped = "Stopped"
	EcsInstanceStatusDeleted = "Deleted"
)

// create -> start -> stop -> delete
func TestEcsInstance(t *testing.T) {

	// init client
	config := getConfigFromEnv()
	ecsClient, err := ecs.NewClientWithAccessKey("cn-hangzhou", config.AccessKeyId, config.AccessKeySecret)
	assertErrorNil(t, err, "Failed to init client")
	fmt.Printf("Init client success\n")

	// get demo instance attributes
	param := getDemoEcsInstanceAttributes(t, ecsClient)

	// create
	instanceId := createEcsInstance(t, ecsClient, param)

	// defer wait for deleted
	defer waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusDeleted, 120)

	// defer delete
	defer deleteEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusStopped, 60)

	// start
	startEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusRunning, 240)

	// stop
	stopEcsInstance(t, ecsClient, instanceId)

	// wait
	waitForEcsInstance(t, ecsClient, instanceId, EcsInstanceStatusStopped, 600)

	//delete all test instance
	deleteAllTestEcsInstance(t, ecsClient)
}

func getDemoEcsInstanceAttributes(t *testing.T, client *ecs.Client) *ecs.DescribeInstanceAttributeResponse {
	fmt.Print("trying to get demo ecs instance...")
	request := ecs.CreateDescribeInstanceAttributeRequest()
	request.InstanceId = getEcsDemoInstanceId()
	response, err := client.DescribeInstanceAttribute(request)
	assertErrorNil(t, err, "Failed to get demo ecs instance attributes")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
	return response
}

func createEcsInstance(t *testing.T, client *ecs.Client, param *ecs.DescribeInstanceAttributeResponse) (instanceId string) {
	fmt.Print("creating ecs instance...")
	request := ecs.CreateCreateInstanceRequest()
	request.ImageId = param.ImageId
	request.InstanceName = InstanceNamePrefix + strconv.FormatInt(time.Now().Unix(), 10)
	request.SecurityGroupId = param.SecurityGroupIds.SecurityGroupId[0]
	request.InstanceType = "ecs.t1.small"
	response, err := client.CreateInstance(request)
	assertErrorNil(t, err, "Failed to create ecs instance")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	instanceId = response.InstanceId
	fmt.Printf("success(%d)! instanceId = %s\n", response.GetHttpStatus(), instanceId)
	return
}

func startEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("starting ecs instance(%s)...", instanceId)
	request := ecs.CreateStartInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.StartInstance(request)
	assertErrorNil(t, err, "Failed to start ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func stopEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("stopping ecs instance(%s)...", instanceId)
	request := ecs.CreateStopInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.StopInstance(request)
	assertErrorNil(t, err, "Failed to stop ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func deleteEcsInstance(t *testing.T, client *ecs.Client, instanceId string) {
	fmt.Printf("deleting ecs instance(%s)...", instanceId)
	request := ecs.CreateDeleteInstanceRequest()
	request.InstanceId = instanceId
	response, err := client.DeleteInstance(request)
	assertErrorNil(t, err, "Failed to delete ecs instance "+instanceId)
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Println("success!")
}

func deleteAllTestEcsInstance(t *testing.T, client *ecs.Client) {
	fmt.Print("list all ecs instances...")
	request := ecs.CreateDescribeInstancesRequest()
	request.PageSize = requests.NewInteger(30)
	request.PageNumber = requests.NewInteger(1)
	response, err := client.DescribeInstances(request)
	assertErrorNil(t, err, "Failed to list all ecs instances ")
	assert.Equal(t, 200, response.GetHttpStatus(), response.GetHttpContentString())
	fmt.Printf("success! TotalCount = %s\n", strconv.Itoa(response.TotalCount))

	for _, instanceInfo := range response.Instances.Instance {
		if strings.HasPrefix(instanceInfo.InstanceName, InstanceNamePrefix) {
			createTime, err := strconv.ParseInt(instanceInfo.InstanceName[26:len(instanceInfo.InstanceName)], 10, 64)
			assertErrorNil(t, err, "Parse instance create time failed: "+instanceInfo.InstanceName)
			if (time.Now().Unix() - createTime) < (60 * 60) {
				fmt.Printf("found undeleted ecs instance(%s) but created in 60 minutes, try to delete next time\n", instanceInfo.InstanceName)
				continue
			} else {
				fmt.Printf("found undeleted ecs instance(%s), status=%s, try to delete it.\n",
					instanceInfo.Status, instanceInfo.InstanceId)
				if instanceInfo.Status == EcsInstanceStatusRunning {
					// stop
					stopEcsInstance(t, client, instanceInfo.InstanceId)
				} else if instanceInfo.Status == EcsInstanceStatusStopped {
					// delete
					deleteEcsInstance(t, client, instanceInfo.InstanceId)
					// wait
					waitForEcsInstance(t, client, instanceInfo.InstanceId, EcsInstanceStatusDeleted, 600)
				}
			}
		}
	}
}

func waitForEcsInstance(t *testing.T, client *ecs.Client, instanceId string, targetStatus string, timeout int) {
	if timeout <= 0 {
		timeout = EcsInstanceDefaultTimeout
	}
	for {
		request := ecs.CreateDescribeInstanceAttributeRequest()
		request.InstanceId = instanceId
		response, err := client.DescribeInstanceAttribute(request)

		if targetStatus == EcsInstanceStatusDeleted {
			if response.GetHttpStatus() == 404 || response.Status == EcsInstanceStatusDeleted {
				fmt.Printf("delete ecs instance(%s) success\n", instanceId)
				break
			} else {
				assertErrorNil(t, err, "Failed to describe ecs instance \n")
			}
		} else {
			assertErrorNil(t, err, "Failed to describe ecs instance \n")
			if response.Status == targetStatus {
				fmt.Printf("ecs instance(%s) status changed to %s, wait a moment\n", instanceId, targetStatus)
				time.Sleep(EcsDefaultWaitForInterval * time.Second)
				break
			} else {
				fmt.Printf("ecs instance(%s) status is %s, wait for changing to %s\n", instanceId, response.Status, targetStatus)
			}
		}

		timeout = timeout - EcsDefaultWaitForInterval
		if timeout <= 0 {
			t.Errorf(fmt.Sprintf("wait for ecs instance(%s) status to %s timeout(%d)\n", instanceId, targetStatus, timeout))
		}
		time.Sleep(EcsDefaultWaitForInterval * time.Second)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"time"

	"github.com/onsi/gomega"

	"github.com/onsi/ginkgo"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/klog"

	"k8s.io/apiserver/pkg/util/logs"
	"k8s.io/kubernetes/pkg/kubectl/util/templates"
	e2e "k8s.io/kubernetes/test/e2e/framework"

	"github.com/openshift/library-go/pkg/serviceability"
	"github.com/openshift/origin/pkg/cmd/flagtypes"
	"github.com/openshift/origin/pkg/monitor"
	testginkgo "github.com/openshift/origin/pkg/test/ginkgo"
	exutil "github.com/openshift/origin/test/extended/util"
)

func main() {
	logs.InitLogs()
	defer logs.FlushLogs()

	rand.Seed(time.Now().UTC().UnixNano())

	root := &cobra.Command{
		Long: templates.LongDesc(`
		OpenShift Tests

		This command verifies behavior of an OpenShift cluster by running remote tests against
		the cluster API that exercise functionality. In general these tests may be disruptive
		or require elevated privileges - see the descriptions of each test suite.
		`),
	}
	flagtypes.GLog(root.PersistentFlags())

	root.AddCommand(
		newRunCommand(),
		newRunUpgradeCommand(),
		newRunTestCommand(),
		newRunMonitorCommand(),
	)

	pflag.CommandLine = pflag.NewFlagSet("empty", pflag.ExitOnError)
	flag.CommandLine = flag.NewFlagSet("empty", flag.ExitOnError)
	exutil.InitStandardFlags()

	if err := func() error {
		defer serviceability.Profile(os.Getenv("OPENSHIFT_PROFILE")).Stop()
		return root.Execute()
	}(); err != nil {
		if ex, ok := err.(testginkgo.ExitError); ok {
			os.Exit(ex.Code)
		}
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
}

func newRunMonitorCommand() *cobra.Command {
	monitorOpt := &monitor.Options{
		Out:    os.Stdout,
		ErrOut: os.Stderr,
	}
	cmd := &cobra.Command{
		Use:   "run-monitor",
		Short: "Continuously verify the cluster is functional",
		Long: templates.LongDesc(`
		Run a continuous verification process

		`),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return monitorOpt.Run()
		},
	}
	return cmd
}

func newRunCommand() *cobra.Command {
	opt := &testginkgo.Options{
		Suites: staticSuites,
	}
	cmd := &cobra.Command{
		Use:   "run SUITE",
		Short: "Run a test suite",
		Long: templates.LongDesc(`
		Run a test suite against an OpenShift server

		This command will run one of the following suites against a cluster identified by the current
		KUBECONFIG file. See the suite description for more on what actions the suite will take.

		If you specify the --dry-run argument, the names of each individual test that is part of the
		suite will be printed, one per line. You may filter this list and pass it back to the run
		command with the --file argument. You may also pipe a list of test names, one per line, on
		standard input by passing "-f -".

		`) + testginkgo.SuitesString(opt.Suites, "\n\nAvailable test suites:\n\n"),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return mirrorToFile(opt, func() error {
				if err := initProvider(opt.Provider); err != nil {
					return err
				}
				os.Setenv("TEST_PROVIDER", opt.Provider)
				e2e.AfterReadingAllFlags(exutil.TestContext)
				return opt.Run(args)
			})
		},
	}
	bindOptions(opt, cmd.Flags())
	return cmd
}

func newRunUpgradeCommand() *cobra.Command {
	opt := &testginkgo.Options{Suites: upgradeSuites}
	upgradeOpt := &UpgradeOptions{}

	cmd := &cobra.Command{
		Use:   "run-upgrade SUITE",
		Short: "Run an upgrade suite",
		Long: templates.LongDesc(`
		Run an upgrade test suite against an OpenShift server

		This command will run one of the following suites against a cluster identified by the current
		KUBECONFIG file. See the suite description for more on what actions the suite will take.

		If you specify the --dry-run argument, the actions the suite will take will be printed to the
		output.

		`) + testginkgo.SuitesString(opt.Suites, "\n\nAvailable upgrade suites:\n\n"),

		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return mirrorToFile(opt, func() error {
				if len(upgradeOpt.ToImage) == 0 {
					return fmt.Errorf("--to-image must be specified to run an upgrade test")
				}

				if len(args) > 0 {
					for _, suite := range opt.Suites {
						if suite.Name == args[0] {
							upgradeOpt.Suite = suite.Name
							upgradeOpt.JUnitDir = opt.JUnitDir
							os.Setenv("TEST_UPGRADE", upgradeOpt.ToEnv())
							break
						}
					}
				}

				if err := initProvider(opt.Provider); err != nil {
					return err
				}
				os.Setenv("TEST_PROVIDER", opt.Provider)
				e2e.AfterReadingAllFlags(exutil.TestContext)
				return opt.Run(args)
			})
		},
	}

	bindOptions(opt, cmd.Flags())
	bindUpgradeOptions(upgradeOpt, cmd.Flags())
	return cmd
}

func newRunTestCommand() *cobra.Command {
	testOpt := &testginkgo.TestOptions{
		Out:    os.Stdout,
		ErrOut: os.Stderr,
	}
	cmd := &cobra.Command{
		Use:   "run-test NAME",
		Short: "Run a single test by name",
		Long: templates.LongDesc(`
		Execute a single test

		This executes a single test by name. It is used by the run command during suite execution but may also
		be used to test in isolation while developing new tests.
		`),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := initProvider(os.Getenv("TEST_PROVIDER")); err != nil {
				return err
			}
			if err := initUpgrade(os.Getenv("TEST_UPGRADE")); err != nil {
				return err
			}
			e2e.AfterReadingAllFlags(exutil.TestContext)
			return testOpt.Run(args)
		},
	}
	cmd.Flags().BoolVar(&testOpt.DryRun, "dry-run", testOpt.DryRun, "Print the test to run without executing them.")
	return cmd
}

// mirrorToFile ensures a copy of all output goes to the provided OutFile, including
// any error returned from fn. The function returns fn() or any error encountered while
// attempting to open the file.
func mirrorToFile(opt *testginkgo.Options, fn func() error) error {
	if len(opt.OutFile) == 0 {
		opt.Out, opt.ErrOut = os.Stdout, os.Stderr
		return fn()
	}

	f, err := os.OpenFile(opt.OutFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0640)
	if err != nil {
		return err
	}
	opt.Out = io.MultiWriter(os.Stdout, f)
	opt.ErrOut = io.MultiWriter(os.Stderr, f)
	exitErr := fn()
	if exitErr != nil {
		fmt.Fprintf(f, "error: %s", exitErr)
	}
	if err := f.Close(); err != nil {
		fmt.Fprintf(os.Stderr, "error: Unable to close output file\n")
	}
	return exitErr
}

func bindOptions(opt *testginkgo.Options, flags *pflag.FlagSet) {
	flags.BoolVar(&opt.DryRun, "dry-run", opt.DryRun, "Print the tests to run without executing them.")
	flags.StringVar(&opt.JUnitDir, "junit-dir", opt.JUnitDir, "The directory to write test reports to.")
	flags.StringVar(&opt.Provider, "provider", opt.Provider, "The cluster infrastructure provider. Will automatically default to the correct value.")
	flags.StringVarP(&opt.TestFile, "file", "f", opt.TestFile, "Create a suite from the newline-delimited test names in this file.")
	flags.StringVarP(&opt.OutFile, "output-file", "o", opt.OutFile, "Write all test output to this file.")
	flags.DurationVar(&opt.Timeout, "timeout", opt.Timeout, "Set the maximum time a test can run before being aborted. This is read from the suite by default, but will be 10 minutes otherwise.")
	flags.BoolVar(&opt.IncludeSuccessOutput, "include-success", opt.IncludeSuccessOutput, "Print output from successful tests.")
}

func initProvider(provider string) error {
	// record the exit error to the output file
	if err := decodeProviderTo(provider, exutil.TestContext); err != nil {
		return err
	}
	exutil.TestContext.AllowedNotReadyNodes = 100
	exutil.TestContext.MaxNodesToGather = 0
	exutil.TestContext.Viper = os.Getenv("VIPERCONFIG")

	// set defaults so these tests don't log
	exutil.TestContext.LoggingSoak.Scale = 1
	exutil.TestContext.LoggingSoak.MilliSecondsBetweenWaves = 5000

	exutil.AnnotateTestSuite()
	exutil.InitTest()
	gomega.RegisterFailHandler(ginkgo.Fail)

	// TODO: infer SSH keys from the cluster
	return nil
}

func decodeProviderTo(provider string, testContext *e2e.TestContextType) error {
	switch provider {
	case "":
		if _, ok := os.LookupEnv("KUBE_SSH_USER"); ok {
			if _, ok := os.LookupEnv("LOCAL_SSH_KEY"); ok {
				testContext.Provider = "local"
			}
		}
		// TODO: detect which provider the cluster is running and use that as a default.
	default:
		var providerInfo struct{ Type string }
		if err := json.Unmarshal([]byte(provider), &providerInfo); err != nil {
			return fmt.Errorf("provider must be a JSON object with the 'type' key at a minimum: %v", err)
		}
		if len(providerInfo.Type) == 0 {
			return fmt.Errorf("provider must be a JSON object with the 'type' key")
		}
		testContext.Provider = providerInfo.Type
		if err := json.Unmarshal([]byte(provider), &testContext.CloudConfig); err != nil {
			return fmt.Errorf("provider must decode into the cloud config object: %v", err)
		}
	}
	if len(testContext.Provider) == 0 {
		testContext.Provider = "skeleton"
	}
	klog.V(2).Infof("Provider %s: %#v", testContext.Provider, testContext.CloudConfig)
	return nil
}
<commit_msg>always dump logs<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"time"

	"github.com/onsi/gomega"

	"github.com/onsi/ginkgo"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/klog"

	"k8s.io/apiserver/pkg/util/logs"
	"k8s.io/kubernetes/pkg/kubectl/util/templates"
	e2e "k8s.io/kubernetes/test/e2e/framework"

	"github.com/openshift/library-go/pkg/serviceability"
	"github.com/openshift/origin/pkg/cmd/flagtypes"
	"github.com/openshift/origin/pkg/monitor"
	testginkgo "github.com/openshift/origin/pkg/test/ginkgo"
	exutil "github.com/openshift/origin/test/extended/util"
)

func main() {
	logs.InitLogs()
	defer logs.FlushLogs()

	rand.Seed(time.Now().UTC().UnixNano())

	root := &cobra.Command{
		Long: templates.LongDesc(`
		OpenShift Tests

		This command verifies behavior of an OpenShift cluster by running remote tests against
		the cluster API that exercise functionality. In general these tests may be disruptive
		or require elevated privileges - see the descriptions of each test suite.
		`),
	}
	flagtypes.GLog(root.PersistentFlags())

	root.AddCommand(
		newRunCommand(),
		newRunUpgradeCommand(),
		newRunTestCommand(),
		newRunMonitorCommand(),
	)

	pflag.CommandLine = pflag.NewFlagSet("empty", pflag.ExitOnError)
	flag.CommandLine = flag.NewFlagSet("empty", flag.ExitOnError)
	exutil.InitStandardFlags()

	if err := func() error {
		defer serviceability.Profile(os.Getenv("OPENSHIFT_PROFILE")).Stop()
		return root.Execute()
	}(); err != nil {
		if ex, ok := err.(testginkgo.ExitError); ok {
			os.Exit(ex.Code)
		}
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
}

func newRunMonitorCommand() *cobra.Command {
	monitorOpt := &monitor.Options{
		Out:    os.Stdout,
		ErrOut: os.Stderr,
	}
	cmd := &cobra.Command{
		Use:   "run-monitor",
		Short: "Continuously verify the cluster is functional",
		Long: templates.LongDesc(`
		Run a continuous verification process

		`),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return monitorOpt.Run()
		},
	}
	return cmd
}

func newRunCommand() *cobra.Command {
	opt := &testginkgo.Options{
		Suites: staticSuites,
	}
	cmd := &cobra.Command{
		Use:   "run SUITE",
		Short: "Run a test suite",
		Long: templates.LongDesc(`
		Run a test suite against an OpenShift server

		This command will run one of the following suites against a cluster identified by the current
		KUBECONFIG file. See the suite description for more on what actions the suite will take.

		If you specify the --dry-run argument, the names of each individual test that is part of the
		suite will be printed, one per line. You may filter this list and pass it back to the run
		command with the --file argument. You may also pipe a list of test names, one per line, on
		standard input by passing "-f -".

		`) + testginkgo.SuitesString(opt.Suites, "\n\nAvailable test suites:\n\n"),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return mirrorToFile(opt, func() error {
				if err := initProvider(opt.Provider); err != nil {
					return err
				}
				os.Setenv("TEST_PROVIDER", opt.Provider)
				e2e.AfterReadingAllFlags(exutil.TestContext)
				exutil.TestContext.DumpLogsOnFailure = true
				return opt.Run(args)
			})
		},
	}
	bindOptions(opt, cmd.Flags())
	return cmd
}

func newRunUpgradeCommand() *cobra.Command {
	opt := &testginkgo.Options{Suites: upgradeSuites}
	upgradeOpt := &UpgradeOptions{}

	cmd := &cobra.Command{
		Use:   "run-upgrade SUITE",
		Short: "Run an upgrade suite",
		Long: templates.LongDesc(`
		Run an upgrade test suite against an OpenShift server

		This command will run one of the following suites against a cluster identified by the current
		KUBECONFIG file. See the suite description for more on what actions the suite will take.

		If you specify the --dry-run argument, the actions the suite will take will be printed to the
		output.

		`) + testginkgo.SuitesString(opt.Suites, "\n\nAvailable upgrade suites:\n\n"),

		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			return mirrorToFile(opt, func() error {
				if len(upgradeOpt.ToImage) == 0 {
					return fmt.Errorf("--to-image must be specified to run an upgrade test")
				}

				if len(args) > 0 {
					for _, suite := range opt.Suites {
						if suite.Name == args[0] {
							upgradeOpt.Suite = suite.Name
							upgradeOpt.JUnitDir = opt.JUnitDir
							os.Setenv("TEST_UPGRADE", upgradeOpt.ToEnv())
							break
						}
					}
				}

				if err := initProvider(opt.Provider); err != nil {
					return err
				}
				os.Setenv("TEST_PROVIDER", opt.Provider)
				e2e.AfterReadingAllFlags(exutil.TestContext)
				exutil.TestContext.DumpLogsOnFailure = true
				return opt.Run(args)
			})
		},
	}

	bindOptions(opt, cmd.Flags())
	bindUpgradeOptions(upgradeOpt, cmd.Flags())
	return cmd
}

func newRunTestCommand() *cobra.Command {
	testOpt := &testginkgo.TestOptions{
		Out:    os.Stdout,
		ErrOut: os.Stderr,
	}
	cmd := &cobra.Command{
		Use:   "run-test NAME",
		Short: "Run a single test by name",
		Long: templates.LongDesc(`
		Execute a single test

		This executes a single test by name. It is used by the run command during suite execution but may also
		be used to test in isolation while developing new tests.
		`),

		SilenceUsage:  true,
		SilenceErrors: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := initProvider(os.Getenv("TEST_PROVIDER")); err != nil {
				return err
			}
			if err := initUpgrade(os.Getenv("TEST_UPGRADE")); err != nil {
				return err
			}
			e2e.AfterReadingAllFlags(exutil.TestContext)
			exutil.TestContext.DumpLogsOnFailure = true
			return testOpt.Run(args)
		},
	}
	cmd.Flags().BoolVar(&testOpt.DryRun, "dry-run", testOpt.DryRun, "Print the test to run without executing them.")
	return cmd
}

// mirrorToFile ensures a copy of all output goes to the provided OutFile, including
// any error returned from fn. The function returns fn() or any error encountered while
// attempting to open the file.
func mirrorToFile(opt *testginkgo.Options, fn func() error) error {
	if len(opt.OutFile) == 0 {
		opt.Out, opt.ErrOut = os.Stdout, os.Stderr
		return fn()
	}

	f, err := os.OpenFile(opt.OutFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0640)
	if err != nil {
		return err
	}
	opt.Out = io.MultiWriter(os.Stdout, f)
	opt.ErrOut = io.MultiWriter(os.Stderr, f)
	exitErr := fn()
	if exitErr != nil {
		fmt.Fprintf(f, "error: %s", exitErr)
	}
	if err := f.Close(); err != nil {
		fmt.Fprintf(os.Stderr, "error: Unable to close output file\n")
	}
	return exitErr
}

func bindOptions(opt *testginkgo.Options, flags *pflag.FlagSet) {
	flags.BoolVar(&opt.DryRun, "dry-run", opt.DryRun, "Print the tests to run without executing them.")
	flags.StringVar(&opt.JUnitDir, "junit-dir", opt.JUnitDir, "The directory to write test reports to.")
	flags.StringVar(&opt.Provider, "provider", opt.Provider, "The cluster infrastructure provider. Will automatically default to the correct value.")
	flags.StringVarP(&opt.TestFile, "file", "f", opt.TestFile, "Create a suite from the newline-delimited test names in this file.")
	flags.StringVarP(&opt.OutFile, "output-file", "o", opt.OutFile, "Write all test output to this file.")
	flags.DurationVar(&opt.Timeout, "timeout", opt.Timeout, "Set the maximum time a test can run before being aborted. This is read from the suite by default, but will be 10 minutes otherwise.")
	flags.BoolVar(&opt.IncludeSuccessOutput, "include-success", opt.IncludeSuccessOutput, "Print output from successful tests.")
}

func initProvider(provider string) error {
	// record the exit error to the output file
	if err := decodeProviderTo(provider, exutil.TestContext); err != nil {
		return err
	}
	exutil.TestContext.AllowedNotReadyNodes = 100
	exutil.TestContext.MaxNodesToGather = 0
	exutil.TestContext.Viper = os.Getenv("VIPERCONFIG")

	// set defaults so these tests don't log
	exutil.TestContext.LoggingSoak.Scale = 1
	exutil.TestContext.LoggingSoak.MilliSecondsBetweenWaves = 5000

	exutil.AnnotateTestSuite()
	exutil.InitTest()
	gomega.RegisterFailHandler(ginkgo.Fail)

	// TODO: infer SSH keys from the cluster
	return nil
}

func decodeProviderTo(provider string, testContext *e2e.TestContextType) error {
	switch provider {
	case "":
		if _, ok := os.LookupEnv("KUBE_SSH_USER"); ok {
			if _, ok := os.LookupEnv("LOCAL_SSH_KEY"); ok {
				testContext.Provider = "local"
			}
		}
		// TODO: detect which provider the cluster is running and use that as a default.
	default:
		var providerInfo struct{ Type string }
		if err := json.Unmarshal([]byte(provider), &providerInfo); err != nil {
			return fmt.Errorf("provider must be a JSON object with the 'type' key at a minimum: %v", err)
		}
		if len(providerInfo.Type) == 0 {
			return fmt.Errorf("provider must be a JSON object with the 'type' key")
		}
		testContext.Provider = providerInfo.Type
		if err := json.Unmarshal([]byte(provider), &testContext.CloudConfig); err != nil {
			return fmt.Errorf("provider must decode into the cloud config object: %v", err)
		}
	}
	if len(testContext.Provider) == 0 {
		testContext.Provider = "skeleton"
	}
	klog.V(2).Infof("Provider %s: %#v", testContext.Provider, testContext.CloudConfig)
	return nil
}
<|endoftext|>


<commit_before>package wtwire

import "io"

// CreateSessionCode is an error code returned by a watchtower in response to a
// CreateSession message. The code directs the client in interpreting the payload
// in the reply.
type CreateSessionCode = ErrorCode

const (
	// CreateSessionCodeAlreadyExists is returned when a session is already
	// active for the public key used to connect to the watchtower. The
	// response includes the serialized reward address in case the original
	// reply was never received and/or processed by the client.
	CreateSessionCodeAlreadyExists CreateSessionCode = 60

	// CreateSessionCodeRejectRejectMaxUpdates the tower rejected the maximum
	// number of state updates proposed by the client.
	CreateSessionCodeRejectRejectMaxUpdates CreateSessionCode = 61

	// CreateSessionCodeRejectRewardRate the tower rejected the reward rate
	// proposed by the client.
	CreateSessionCodeRejectRewardRate CreateSessionCode = 62

	// CreateSessionCodeRejectSweepFeeRate the tower rejected the sweep fee
	// rate proposed by the client.
	CreateSessionCodeRejectSweepFeeRate CreateSessionCode = 63

	// CreateSessionCodeRejectBlobType is returned when the tower does not
	// support the proposed blob type.
	CreateSessionCodeRejectBlobType CreateSessionCode = 64
)

// MaxCreateSessionReplyDataLength is the maximum size of the Data payload
// returned in a CreateSessionReply message. This does not include the length of
// the Data field, which is a varint up to 3 bytes in size.
const MaxCreateSessionReplyDataLength = 1024

// CreateSessionReply is a message sent from watchtower to client in response to a
// CreateSession message, and signals either an acceptance or rejection of the
// proposed session parameters.
type CreateSessionReply struct {
	// Code will be non-zero if the watchtower rejected the session init.
	Code CreateSessionCode

	// Data is a byte slice returned the caller of the message, and is to be
	// interpreted according to the error Code. When the response is
	// CreateSessionCodeOK, data encodes the reward address to be included in
	// any sweep transactions if the reward is not dusty. Otherwise, it may
	// encode the watchtowers configured parameters for any policy
	// rejections.
	Data []byte
}

// A compile time check to ensure CreateSessionReply implements the wtwire.Message
// interface.
var _ Message = (*CreateSessionReply)(nil)

// Decode deserializes a serialized CreateSessionReply message stored in the passed
// io.Reader observing the specified protocol version.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) Decode(r io.Reader, pver uint32) error {
	return ReadElements(r,
		&m.Code,
		&m.Data,
	)
}

// Encode serializes the target CreateSessionReply into the passed io.Writer
// observing the protocol version specified.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) Encode(w io.Writer, pver uint32) error {
	return WriteElements(w,
		m.Code,
		m.Data,
	)
}

// MsgType returns the integer uniquely identifying this message type on the
// wire.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) MsgType() MessageType {
	return MsgCreateSessionReply
}

// MaxPayloadLength returns the maximum allowed payload size for a CreateSessionReply
// complete message observing the specified protocol version.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) MaxPayloadLength(uint32) uint32 {
	return 2 + 3 + MaxCreateSessionReplyDataLength
}
<commit_msg>watchtower/wtwire/create_session_reply: remove extra Reject from code<commit_after>package wtwire

import "io"

// CreateSessionCode is an error code returned by a watchtower in response to a
// CreateSession message. The code directs the client in interpreting the payload
// in the reply.
type CreateSessionCode = ErrorCode

const (
	// CreateSessionCodeAlreadyExists is returned when a session is already
	// active for the public key used to connect to the watchtower. The
	// response includes the serialized reward address in case the original
	// reply was never received and/or processed by the client.
	CreateSessionCodeAlreadyExists CreateSessionCode = 60

	// CreateSessionCodeRejectMaxUpdates the tower rejected the maximum
	// number of state updates proposed by the client.
	CreateSessionCodeRejectMaxUpdates CreateSessionCode = 61

	// CreateSessionCodeRejectRewardRate the tower rejected the reward rate
	// proposed by the client.
	CreateSessionCodeRejectRewardRate CreateSessionCode = 62

	// CreateSessionCodeRejectSweepFeeRate the tower rejected the sweep fee
	// rate proposed by the client.
	CreateSessionCodeRejectSweepFeeRate CreateSessionCode = 63

	// CreateSessionCodeRejectBlobType is returned when the tower does not
	// support the proposed blob type.
	CreateSessionCodeRejectBlobType CreateSessionCode = 64
)

// MaxCreateSessionReplyDataLength is the maximum size of the Data payload
// returned in a CreateSessionReply message. This does not include the length of
// the Data field, which is a varint up to 3 bytes in size.
const MaxCreateSessionReplyDataLength = 1024

// CreateSessionReply is a message sent from watchtower to client in response to a
// CreateSession message, and signals either an acceptance or rejection of the
// proposed session parameters.
type CreateSessionReply struct {
	// Code will be non-zero if the watchtower rejected the session init.
	Code CreateSessionCode

	// Data is a byte slice returned the caller of the message, and is to be
	// interpreted according to the error Code. When the response is
	// CreateSessionCodeOK, data encodes the reward address to be included in
	// any sweep transactions if the reward is not dusty. Otherwise, it may
	// encode the watchtowers configured parameters for any policy
	// rejections.
	Data []byte
}

// A compile time check to ensure CreateSessionReply implements the wtwire.Message
// interface.
var _ Message = (*CreateSessionReply)(nil)

// Decode deserializes a serialized CreateSessionReply message stored in the passed
// io.Reader observing the specified protocol version.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) Decode(r io.Reader, pver uint32) error {
	return ReadElements(r,
		&m.Code,
		&m.Data,
	)
}

// Encode serializes the target CreateSessionReply into the passed io.Writer
// observing the protocol version specified.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) Encode(w io.Writer, pver uint32) error {
	return WriteElements(w,
		m.Code,
		m.Data,
	)
}

// MsgType returns the integer uniquely identifying this message type on the
// wire.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) MsgType() MessageType {
	return MsgCreateSessionReply
}

// MaxPayloadLength returns the maximum allowed payload size for a CreateSessionReply
// complete message observing the specified protocol version.
//
// This is part of the wtwire.Message interface.
func (m *CreateSessionReply) MaxPayloadLength(uint32) uint32 {
	return 2 + 3 + MaxCreateSessionReplyDataLength
}
<|endoftext|>


<commit_before>package synctest

import (
	"context"
	"fmt"
	"time"

	"koding/klient/machine/index"
	"koding/klient/machine/index/indextest"
	msync "koding/klient/machine/mount/sync"
)

// SyncLocal sends all index changes between local and remote directories to
// provided syncer. Then it calls produced Execers. Returned contex will be
// closed when all changes are consumed and their events executed. The direction
// of changes is marked by dir argument and it must be ChangeMetaRemote and/or
// ChangeMetaRemote
func SyncLocal(s msync.Syncer, rootA, rootB string, dir index.ChangeMeta) (context.Context, context.CancelFunc, error) {
	if dir&^(index.ChangeMetaLocal|index.ChangeMetaRemote) != 0 {
		return nil, nil, fmt.Errorf("invalid change type: %v", dir)
	}

	cs, err := indextest.Compare(rootA, rootB)
	if err != nil {
		return nil, nil, err
	}

	evs := make([]*msync.Event, 0, len(cs))
	for _, c := range cs {
		c = index.NewChange(c.Path(), c.Meta()|dir)
		evs = append(evs, msync.NewEvent(context.Background(), nil, c))
	}

	var (
		evC = make(chan *msync.Event)
		exC = s.ExecStream(evC)
	)

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		defer cancel()
		for i := 0; i < len(evs); {
			select {
			case evC <- evs[i]:
			case ex := <-exC:
				ex.Exec()
				i++
			case <-ctx.Done():
				return
			}
		}
	}()

	return ctx, cancel, nil
}

// ExecChange uses provided to create Execer from provided index change. Then
// the function will execute it end return its error or timeout if such occurs.
func ExecChange(s msync.Syncer, change *index.Change, timeout time.Duration) error {
	var (
		ev       = msync.NewEvent(context.Background(), nil, change)
		evC      = make(chan *msync.Event)
		timeoutC = time.After(timeout)
	)

	go func() {
		select {
		case evC <- ev:
		case <-timeoutC:
		}
	}()

	select {
	case ex := <-s.ExecStream(evC):
		if ex == nil {
			return fmt.Errorf("nil execer received")
		}
		return ex.Exec()
	case <-timeoutC:
		return fmt.Errorf("timed out after %s", timeout)
	}
}
<commit_msg>klient/machine: add synctest package that helps testing sync interface<commit_after>package synctest

import (
	"context"
	"fmt"
	"time"

	"koding/klient/machine/index"
	"koding/klient/machine/index/indextest"
	msync "koding/klient/machine/mount/sync"
)

// SyncLocal sends all index changes between local and remote directories to
// provided syncer. Then it calls produced Execers. Returned contex will be
// closed when all changes are consumed and their events executed. The direction
// of changes is marked by dir argument and it must be ChangeMetaRemote and/or
// ChangeMetaRemote.
func SyncLocal(s msync.Syncer, rootA, rootB string, dir index.ChangeMeta) (context.Context, context.CancelFunc, error) {
	if dir&^(index.ChangeMetaLocal|index.ChangeMetaRemote) != 0 {
		return nil, nil, fmt.Errorf("invalid change type: %v", dir)
	}

	cs, err := indextest.Compare(rootA, rootB)
	if err != nil {
		return nil, nil, err
	}

	evs := make([]*msync.Event, 0, len(cs))
	for _, c := range cs {
		c = index.NewChange(c.Path(), c.Meta()|dir)
		evs = append(evs, msync.NewEvent(context.Background(), nil, c))
	}

	var (
		evC = make(chan *msync.Event)
		exC = s.ExecStream(evC)
	)

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		defer cancel()
		for i := 0; i < len(evs); {
			select {
			case evC <- evs[i]:
			case ex := <-exC:
				ex.Exec()
				i++
			case <-ctx.Done():
				return
			}
		}
	}()

	return ctx, cancel, nil
}

// ExecChange uses provided to create Execer from provided index change. Then
// the function will execute it end return its error or timeout if such occurs.
func ExecChange(s msync.Syncer, change *index.Change, timeout time.Duration) error {
	var (
		ev       = msync.NewEvent(context.Background(), nil, change)
		evC      = make(chan *msync.Event)
		timeoutC = time.After(timeout)
	)

	go func() {
		select {
		case evC <- ev:
		case <-timeoutC:
		}
	}()

	select {
	case ex := <-s.ExecStream(evC):
		if ex == nil {
			return fmt.Errorf("nil execer received")
		}
		return ex.Exec()
	case <-timeoutC:
		return fmt.Errorf("timed out after %s", timeout)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 - 2017 Ka-Hing Cheung
//
// 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 internal

import (
	"io"
	"runtime"
	"runtime/debug"
	"sync"

	"github.com/jacobsa/fuse"
	"github.com/shirou/gopsutil/mem"
)

type BufferPool struct {
	mu   sync.Mutex
	cond *sync.Cond

	numBuffers uint64
	maxBuffers uint64

	totalBuffers       uint64
	computedMaxbuffers uint64

	pool *sync.Pool
}

const BUF_SIZE = 5 * 1024 * 1024

func maxMemToUse(buffersNow uint64) uint64 {
	m, err := mem.VirtualMemory()
	if err != nil {
		panic(err)
	}

	availableMem, err := getCgroupAvailableMem()
	log.Debugf("amount of available memory from cgroup is: %v", availableMem/1024/1024)

	if err != nil || availableMem < 0 || availableMem > m.Available {
		availableMem = m.Available
	}

	log.Debugf("amount of available memory: %v", availableMem/1024/1024)

	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)

	log.Debugf("amount of allocated memory: %v %v", ms.Sys/1024/1024, ms.Alloc/1024/1024)

	max := uint64(availableMem+ms.Sys) / 2
	maxbuffers := MaxUInt64(max/BUF_SIZE, 1)
	log.Debugf("using up to %v %vMB buffers, now is %v", maxbuffers, BUF_SIZE/1024/1024, buffersNow)
	return maxbuffers
}

func rounduUp(size uint64, pageSize int) int {
	return pages(size, pageSize) * pageSize
}

func pages(size uint64, pageSize int) int {
	return int((size + uint64(pageSize) - 1) / uint64(pageSize))
}

func (pool BufferPool) Init() *BufferPool {
	pool.cond = sync.NewCond(&pool.mu)

	pool.computedMaxbuffers = pool.maxBuffers
	pool.pool = &sync.Pool{New: func() interface{} {
		return make([]byte, 0, BUF_SIZE)
	}}

	return &pool
}

// for testing
func NewBufferPool(maxSizeGlobal uint64) *BufferPool {
	pool := BufferPool{maxBuffers: maxSizeGlobal / BUF_SIZE}.Init()
	return pool
}

func (pool *BufferPool) RequestBuffer() (buf []byte) {
	return pool.RequestMultiple(BUF_SIZE, true)[0]
}

func (pool *BufferPool) recomputeBufferLimit() {
	if pool.maxBuffers == 0 {
		pool.computedMaxbuffers = maxMemToUse(pool.numBuffers)
		if pool.computedMaxbuffers == 0 {
			panic("OOM")
		}
	}
}

func (pool *BufferPool) RequestMultiple(size uint64, block bool) (buffers [][]byte) {
	nPages := pages(size, BUF_SIZE)

	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.totalBuffers%10 == 0 {
		pool.recomputeBufferLimit()
	}

	for pool.numBuffers+uint64(nPages) > pool.computedMaxbuffers {
		if block {
			if pool.numBuffers == 0 {
				pool.MaybeGC()
				pool.recomputeBufferLimit()
				if pool.numBuffers+uint64(nPages) > pool.computedMaxbuffers {
					// we don't have any in use buffers, and we've made attempts to
					// free memory AND correct our limits, yet we still can't allocate.
					// it's likely that we are simply asking for too much
					log.Errorf("Unable to allocate %d bytes, limit is %d bytes",
						nPages*BUF_SIZE, pool.computedMaxbuffers*BUF_SIZE)
					panic("OOM")
				}
			}
			pool.cond.Wait()
		} else {
			return
		}
	}

	for i := 0; i < nPages; i++ {
		pool.numBuffers++
		pool.totalBuffers++
		buf := pool.pool.Get()
		buffers = append(buffers, buf.([]byte))
	}
	return
}

func (pool *BufferPool) MaybeGC() {
	if pool.numBuffers == 0 {
		debug.FreeOSMemory()
	}
}

func (pool *BufferPool) Free(buf []byte) {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	buf = buf[:0]
	pool.pool.Put(buf)
	pool.numBuffers--
	pool.cond.Signal()
}

var mbufLog = GetLogger("mbuf")

type MBuf struct {
	pool    *BufferPool
	buffers [][]byte
	rbuf    int
	wbuf    int
	rp      int
	wp      int
}

func (mb MBuf) Init(h *BufferPool, size uint64, block bool) *MBuf {
	mb.pool = h

	if size != 0 {
		mb.buffers = h.RequestMultiple(size, block)
		if mb.buffers == nil {
			return nil
		}
	}

	return &mb
}

func (mb *MBuf) Len() (length int) {
	for i := mb.rbuf; i < int(len(mb.buffers)); i++ {
		var bufSize int
		var start int
		if i == mb.wbuf {
			bufSize = mb.wp
		} else {
			bufSize = int(len(mb.buffers[i]))
		}

		if i == mb.rbuf {
			start = mb.rp
		} else {
			start = 0
		}

		length += bufSize - start
	}

	return
}

// seek only seeks the reader
func (mb *MBuf) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case 0: // relative to beginning
		if offset == 0 {
			mb.rbuf = 0
			mb.rp = 0
			return 0, nil
		}
	case 1: // relative to current position
		if offset == 0 {
			for i := 0; i < mb.rbuf; i++ {
				offset += int64(len(mb.buffers[i]))
			}
			offset += int64(mb.rp)
			return offset, nil
		}

	case 2: // relative to the end
		if offset == 0 {
			for i := 0; i < len(mb.buffers); i++ {
				offset += int64(len(mb.buffers[i]))
			}
			mb.rbuf = len(mb.buffers)
			mb.rp = 0
			return offset, nil
		}
	}

	log.Errorf("Seek %d %d", offset, whence)
	panic(fuse.EINVAL)

	return 0, fuse.EINVAL
}

func (mb *MBuf) Read(p []byte) (n int, err error) {
	if mb.rbuf == mb.wbuf && mb.rp == mb.wp {
		err = io.EOF
		return
	}

	if mb.rp == cap(mb.buffers[mb.rbuf]) {
		mb.rbuf++
		mb.rp = 0
	}

	if mb.rbuf == len(mb.buffers) {
		err = io.EOF
		return
	} else if mb.rbuf > len(mb.buffers) {
		panic("mb.cur > len(mb.buffers)")
	}

	n = copy(p, mb.buffers[mb.rbuf][mb.rp:])
	mb.rp += n

	return
}

func (mb *MBuf) Full() bool {
	return mb.buffers == nil || (mb.wp == cap(mb.buffers[mb.wbuf]) && mb.wbuf+1 == len(mb.buffers))
}

func (mb *MBuf) Write(p []byte) (n int, err error) {
	b := mb.buffers[mb.wbuf]

	if mb.wp == cap(b) {
		if mb.wbuf+1 == len(mb.buffers) {
			return
		}
		mb.wbuf++
		b = mb.buffers[mb.wbuf]
		mb.wp = 0
	} else if mb.wp > cap(b) {
		panic("mb.wp > cap(b)")
	}

	n = copy(b[mb.wp:cap(b)], p)
	mb.wp += n
	// resize the buffer to account for what we just read
	mb.buffers[mb.wbuf] = mb.buffers[mb.wbuf][:mb.wp]

	return
}

func (mb *MBuf) WriteFrom(r io.Reader) (n int, err error) {
	b := mb.buffers[mb.wbuf]

	if mb.wp == cap(b) {
		if mb.wbuf+1 == len(mb.buffers) {
			return
		}
		mb.wbuf++
		b = mb.buffers[mb.wbuf]
		mb.wp = 0
	} else if mb.wp > cap(b) {
		panic("mb.wp > cap(b)")
	}

	n, err = r.Read(b[mb.wp:cap(b)])
	mb.wp += n
	// resize the buffer to account for what we just read
	mb.buffers[mb.wbuf] = mb.buffers[mb.wbuf][:mb.wp]

	return
}

func (mb *MBuf) Free() {
	for _, b := range mb.buffers {
		mb.pool.Free(b)
	}

	mb.buffers = nil
}

var bufferLog = GetLogger("buffer")

type Buffer struct {
	mu   sync.Mutex
	cond *sync.Cond

	buf    *MBuf
	reader io.ReadCloser
	err    error
}

type ReaderProvider func() (io.ReadCloser, error)

func (b Buffer) Init(buf *MBuf, r ReaderProvider) *Buffer {
	b.buf = buf
	b.cond = sync.NewCond(&b.mu)

	go func() {
		b.readLoop(r)
	}()

	return &b
}

func (b *Buffer) readLoop(r ReaderProvider) {
	for {
		b.mu.Lock()
		if b.reader == nil {
			b.reader, b.err = r()
			b.cond.Broadcast()
			if b.err != nil {
				b.mu.Unlock()
				break
			}
		}

		if b.buf == nil {
			// buffer was drained
			b.mu.Unlock()
			break
		}

		nread, err := b.buf.WriteFrom(b.reader)
		if err != nil {
			b.err = err
			b.mu.Unlock()
			break
		}
		bufferLog.Debugf("wrote %v into buffer", nread)

		if nread == 0 {
			b.reader.Close()
			b.mu.Unlock()
			break
		}

		b.mu.Unlock()
		// if we get here we've read _something_, bounce this goroutine
		// to allow another one to read
		runtime.Gosched()
	}
	bufferLog.Debugf("<-- readLoop()")
}

func (b *Buffer) readFromStream(p []byte) (n int, err error) {
	bufferLog.Debugf("reading %v from stream", len(p))

	n, err = b.reader.Read(p)
	if n != 0 && err == io.ErrUnexpectedEOF {
		err = nil
	} else {
		bufferLog.Debugf("read %v from stream", n)
	}
	return
}

func (b *Buffer) Read(p []byte) (n int, err error) {
	bufferLog.Debugf("Buffer.Read(%v)", len(p))

	b.mu.Lock()
	defer b.mu.Unlock()

	for b.reader == nil && b.err == nil {
		bufferLog.Debugf("waiting for stream")
		b.cond.Wait()
		if b.err != nil {
			err = b.err
			return
		}
	}

	// we could have received the err before Read was called
	if b.reader == nil {
		if b.err == nil {
			panic("reader and err are both nil")
		}
		err = b.err
		return
	}

	if b.buf != nil {
		bufferLog.Debugf("reading %v from buffer", len(p))

		n, err = b.buf.Read(p)
		if n == 0 {
			b.buf.Free()
			b.buf = nil
			bufferLog.Debugf("drained buffer")
			n, err = b.readFromStream(p)
		} else {
			bufferLog.Debugf("read %v from buffer", n)
		}
	} else if b.err != nil {
		err = b.err
	} else {
		n, err = b.readFromStream(p)
	}

	return
}

func (b *Buffer) Close() (err error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.reader != nil {
		err = b.reader.Close()
	}

	if b.buf != nil {
		b.buf.Free()
		b.buf = nil
	}

	return
}
<commit_msg>only log if we can get it from cgroup<commit_after>// Copyright 2015 - 2017 Ka-Hing Cheung
//
// 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 internal

import (
	"io"
	"runtime"
	"runtime/debug"
	"sync"

	"github.com/jacobsa/fuse"
	"github.com/shirou/gopsutil/mem"
)

type BufferPool struct {
	mu   sync.Mutex
	cond *sync.Cond

	numBuffers uint64
	maxBuffers uint64

	totalBuffers       uint64
	computedMaxbuffers uint64

	pool *sync.Pool
}

const BUF_SIZE = 5 * 1024 * 1024

func maxMemToUse(buffersNow uint64) uint64 {
	m, err := mem.VirtualMemory()
	if err != nil {
		panic(err)
	}

	availableMem, err := getCgroupAvailableMem()
	if err != nil {
		log.Debugf("amount of available memory from cgroup is: %v", availableMem/1024/1024)
	}

	if err != nil || availableMem < 0 || availableMem > m.Available {
		availableMem = m.Available
	}

	log.Debugf("amount of available memory: %v", availableMem/1024/1024)

	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)

	log.Debugf("amount of allocated memory: %v %v", ms.Sys/1024/1024, ms.Alloc/1024/1024)

	max := uint64(availableMem+ms.Sys) / 2
	maxbuffers := MaxUInt64(max/BUF_SIZE, 1)
	log.Debugf("using up to %v %vMB buffers, now is %v", maxbuffers, BUF_SIZE/1024/1024, buffersNow)
	return maxbuffers
}

func rounduUp(size uint64, pageSize int) int {
	return pages(size, pageSize) * pageSize
}

func pages(size uint64, pageSize int) int {
	return int((size + uint64(pageSize) - 1) / uint64(pageSize))
}

func (pool BufferPool) Init() *BufferPool {
	pool.cond = sync.NewCond(&pool.mu)

	pool.computedMaxbuffers = pool.maxBuffers
	pool.pool = &sync.Pool{New: func() interface{} {
		return make([]byte, 0, BUF_SIZE)
	}}

	return &pool
}

// for testing
func NewBufferPool(maxSizeGlobal uint64) *BufferPool {
	pool := BufferPool{maxBuffers: maxSizeGlobal / BUF_SIZE}.Init()
	return pool
}

func (pool *BufferPool) RequestBuffer() (buf []byte) {
	return pool.RequestMultiple(BUF_SIZE, true)[0]
}

func (pool *BufferPool) recomputeBufferLimit() {
	if pool.maxBuffers == 0 {
		pool.computedMaxbuffers = maxMemToUse(pool.numBuffers)
		if pool.computedMaxbuffers == 0 {
			panic("OOM")
		}
	}
}

func (pool *BufferPool) RequestMultiple(size uint64, block bool) (buffers [][]byte) {
	nPages := pages(size, BUF_SIZE)

	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.totalBuffers%10 == 0 {
		pool.recomputeBufferLimit()
	}

	for pool.numBuffers+uint64(nPages) > pool.computedMaxbuffers {
		if block {
			if pool.numBuffers == 0 {
				pool.MaybeGC()
				pool.recomputeBufferLimit()
				if pool.numBuffers+uint64(nPages) > pool.computedMaxbuffers {
					// we don't have any in use buffers, and we've made attempts to
					// free memory AND correct our limits, yet we still can't allocate.
					// it's likely that we are simply asking for too much
					log.Errorf("Unable to allocate %d bytes, limit is %d bytes",
						nPages*BUF_SIZE, pool.computedMaxbuffers*BUF_SIZE)
					panic("OOM")
				}
			}
			pool.cond.Wait()
		} else {
			return
		}
	}

	for i := 0; i < nPages; i++ {
		pool.numBuffers++
		pool.totalBuffers++
		buf := pool.pool.Get()
		buffers = append(buffers, buf.([]byte))
	}
	return
}

func (pool *BufferPool) MaybeGC() {
	if pool.numBuffers == 0 {
		debug.FreeOSMemory()
	}
}

func (pool *BufferPool) Free(buf []byte) {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	buf = buf[:0]
	pool.pool.Put(buf)
	pool.numBuffers--
	pool.cond.Signal()
}

var mbufLog = GetLogger("mbuf")

type MBuf struct {
	pool    *BufferPool
	buffers [][]byte
	rbuf    int
	wbuf    int
	rp      int
	wp      int
}

func (mb MBuf) Init(h *BufferPool, size uint64, block bool) *MBuf {
	mb.pool = h

	if size != 0 {
		mb.buffers = h.RequestMultiple(size, block)
		if mb.buffers == nil {
			return nil
		}
	}

	return &mb
}

func (mb *MBuf) Len() (length int) {
	for i := mb.rbuf; i < int(len(mb.buffers)); i++ {
		var bufSize int
		var start int
		if i == mb.wbuf {
			bufSize = mb.wp
		} else {
			bufSize = int(len(mb.buffers[i]))
		}

		if i == mb.rbuf {
			start = mb.rp
		} else {
			start = 0
		}

		length += bufSize - start
	}

	return
}

// seek only seeks the reader
func (mb *MBuf) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case 0: // relative to beginning
		if offset == 0 {
			mb.rbuf = 0
			mb.rp = 0
			return 0, nil
		}
	case 1: // relative to current position
		if offset == 0 {
			for i := 0; i < mb.rbuf; i++ {
				offset += int64(len(mb.buffers[i]))
			}
			offset += int64(mb.rp)
			return offset, nil
		}

	case 2: // relative to the end
		if offset == 0 {
			for i := 0; i < len(mb.buffers); i++ {
				offset += int64(len(mb.buffers[i]))
			}
			mb.rbuf = len(mb.buffers)
			mb.rp = 0
			return offset, nil
		}
	}

	log.Errorf("Seek %d %d", offset, whence)
	panic(fuse.EINVAL)

	return 0, fuse.EINVAL
}

func (mb *MBuf) Read(p []byte) (n int, err error) {
	if mb.rbuf == mb.wbuf && mb.rp == mb.wp {
		err = io.EOF
		return
	}

	if mb.rp == cap(mb.buffers[mb.rbuf]) {
		mb.rbuf++
		mb.rp = 0
	}

	if mb.rbuf == len(mb.buffers) {
		err = io.EOF
		return
	} else if mb.rbuf > len(mb.buffers) {
		panic("mb.cur > len(mb.buffers)")
	}

	n = copy(p, mb.buffers[mb.rbuf][mb.rp:])
	mb.rp += n

	return
}

func (mb *MBuf) Full() bool {
	return mb.buffers == nil || (mb.wp == cap(mb.buffers[mb.wbuf]) && mb.wbuf+1 == len(mb.buffers))
}

func (mb *MBuf) Write(p []byte) (n int, err error) {
	b := mb.buffers[mb.wbuf]

	if mb.wp == cap(b) {
		if mb.wbuf+1 == len(mb.buffers) {
			return
		}
		mb.wbuf++
		b = mb.buffers[mb.wbuf]
		mb.wp = 0
	} else if mb.wp > cap(b) {
		panic("mb.wp > cap(b)")
	}

	n = copy(b[mb.wp:cap(b)], p)
	mb.wp += n
	// resize the buffer to account for what we just read
	mb.buffers[mb.wbuf] = mb.buffers[mb.wbuf][:mb.wp]

	return
}

func (mb *MBuf) WriteFrom(r io.Reader) (n int, err error) {
	b := mb.buffers[mb.wbuf]

	if mb.wp == cap(b) {
		if mb.wbuf+1 == len(mb.buffers) {
			return
		}
		mb.wbuf++
		b = mb.buffers[mb.wbuf]
		mb.wp = 0
	} else if mb.wp > cap(b) {
		panic("mb.wp > cap(b)")
	}

	n, err = r.Read(b[mb.wp:cap(b)])
	mb.wp += n
	// resize the buffer to account for what we just read
	mb.buffers[mb.wbuf] = mb.buffers[mb.wbuf][:mb.wp]

	return
}

func (mb *MBuf) Free() {
	for _, b := range mb.buffers {
		mb.pool.Free(b)
	}

	mb.buffers = nil
}

var bufferLog = GetLogger("buffer")

type Buffer struct {
	mu   sync.Mutex
	cond *sync.Cond

	buf    *MBuf
	reader io.ReadCloser
	err    error
}

type ReaderProvider func() (io.ReadCloser, error)

func (b Buffer) Init(buf *MBuf, r ReaderProvider) *Buffer {
	b.buf = buf
	b.cond = sync.NewCond(&b.mu)

	go func() {
		b.readLoop(r)
	}()

	return &b
}

func (b *Buffer) readLoop(r ReaderProvider) {
	for {
		b.mu.Lock()
		if b.reader == nil {
			b.reader, b.err = r()
			b.cond.Broadcast()
			if b.err != nil {
				b.mu.Unlock()
				break
			}
		}

		if b.buf == nil {
			// buffer was drained
			b.mu.Unlock()
			break
		}

		nread, err := b.buf.WriteFrom(b.reader)
		if err != nil {
			b.err = err
			b.mu.Unlock()
			break
		}
		bufferLog.Debugf("wrote %v into buffer", nread)

		if nread == 0 {
			b.reader.Close()
			b.mu.Unlock()
			break
		}

		b.mu.Unlock()
		// if we get here we've read _something_, bounce this goroutine
		// to allow another one to read
		runtime.Gosched()
	}
	bufferLog.Debugf("<-- readLoop()")
}

func (b *Buffer) readFromStream(p []byte) (n int, err error) {
	bufferLog.Debugf("reading %v from stream", len(p))

	n, err = b.reader.Read(p)
	if n != 0 && err == io.ErrUnexpectedEOF {
		err = nil
	} else {
		bufferLog.Debugf("read %v from stream", n)
	}
	return
}

func (b *Buffer) Read(p []byte) (n int, err error) {
	bufferLog.Debugf("Buffer.Read(%v)", len(p))

	b.mu.Lock()
	defer b.mu.Unlock()

	for b.reader == nil && b.err == nil {
		bufferLog.Debugf("waiting for stream")
		b.cond.Wait()
		if b.err != nil {
			err = b.err
			return
		}
	}

	// we could have received the err before Read was called
	if b.reader == nil {
		if b.err == nil {
			panic("reader and err are both nil")
		}
		err = b.err
		return
	}

	if b.buf != nil {
		bufferLog.Debugf("reading %v from buffer", len(p))

		n, err = b.buf.Read(p)
		if n == 0 {
			b.buf.Free()
			b.buf = nil
			bufferLog.Debugf("drained buffer")
			n, err = b.readFromStream(p)
		} else {
			bufferLog.Debugf("read %v from buffer", n)
		}
	} else if b.err != nil {
		err = b.err
	} else {
		n, err = b.readFromStream(p)
	}

	return
}

func (b *Buffer) Close() (err error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.reader != nil {
		err = b.reader.Close()
	}

	if b.buf != nil {
		b.buf.Free()
		b.buf = nil
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2021 Google 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 jwks

import (
	"time"

	"github.com/google/exposure-notifications-server/internal/setup"
	"github.com/google/exposure-notifications-server/pkg/database"
	"github.com/google/exposure-notifications-server/pkg/secrets"
)

var (
	_ setup.DatabaseConfigProvider      = (*Config)(nil)
	_ setup.SecretManagerConfigProvider = (*Config)(nil)
)

type Config struct {
	Database      database.Config
	SecretManager secrets.Config

	Port string `env:"PORT, default=8080"`

	// MaxRuntime is how long an individual handler should run.
	MaxRuntime time.Duration `env:"MAX_RUNTIME, default=10m"`

	// RequestTimeout is the client per-request timeout when accessing
	// remote JWKS documents.
	RequestTimeout time.Duration `env:"REQUEST_TIMEOUT, default=30s"`

	KeyCleanupTTL time.Duration `env:"HEALTH_AUTHORITY_KEY_CLEANUP_TTL, default=720h"` // 30 days

	// MaxWorkers is the number of parallel JWKS updates that can occur.
	MaxWorkers uint `env:"MAX_WORKERS, default=5"`
}

func (c *Config) DatabaseConfig() *database.Config {
	return &c.Database
}

func (c *Config) SecretManagerConfig() *secrets.Config {
	return &c.SecretManager
}
<commit_msg>Add o11y exporter to jwks (#1486)<commit_after>// Copyright 2021 Google 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 jwks

import (
	"time"

	"github.com/google/exposure-notifications-server/internal/setup"
	"github.com/google/exposure-notifications-server/pkg/database"
	"github.com/google/exposure-notifications-server/pkg/observability"
	"github.com/google/exposure-notifications-server/pkg/secrets"
)

var (
	_ setup.DatabaseConfigProvider              = (*Config)(nil)
	_ setup.ObservabilityExporterConfigProvider = (*Config)(nil)
	_ setup.SecretManagerConfigProvider         = (*Config)(nil)
)

type Config struct {
	Database              database.Config
	ObservabilityExporter observability.Config
	SecretManager         secrets.Config

	Port string `env:"PORT, default=8080"`

	// MaxRuntime is how long an individual handler should run.
	MaxRuntime time.Duration `env:"MAX_RUNTIME, default=10m"`

	// RequestTimeout is the client per-request timeout when accessing
	// remote JWKS documents.
	RequestTimeout time.Duration `env:"REQUEST_TIMEOUT, default=30s"`

	KeyCleanupTTL time.Duration `env:"HEALTH_AUTHORITY_KEY_CLEANUP_TTL, default=720h"` // 30 days

	// MaxWorkers is the number of parallel JWKS updates that can occur.
	MaxWorkers uint `env:"MAX_WORKERS, default=5"`
}

func (c *Config) DatabaseConfig() *database.Config {
	return &c.Database
}

func (c *Config) ObservabilityExporterConfig() *observability.Config {
	return &c.ObservabilityExporter
}

func (c *Config) SecretManagerConfig() *secrets.Config {
	return &c.SecretManager
}
<|endoftext|>


<commit_before>package models

import (
	"database/sql"
	"fmt"
	"log"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/upsert"
	"github.com/jmoiron/sqlx"
)

// Stop is a single transit stop for a particular route. If a
// stop serves more than one route, there are multiple distinct
// entries for that stop.
type Stop struct {
	ID       string `json:"stop_id" db:"stop_id" upsert:"key"`
	RouteID  string `json:"route_id" db:"route_id" upsert:"key"`
	AgencyID string `json:"agency_id" db:"agency_id" upsert:"key"`
	Name     string `json:"stop_name" db:"stop_name"`

	DirectionID int    `json:"direction_id" db:"direction_id"`
	Headsign    string `json:"headsign" db:"headsign"`

	Lat float64 `json:"lat" db:"lat" upsert:"omit"`
	Lon float64 `json:"lon" db:"lon" upsert:"omit"`

	// Location is an "earth" field value that combines lat and lon into
	// a single field.
	Location interface{} `json:"-" db:"location" upsert_value:"ll_to_earth(:lat, :lon)"`

	// StopSequence is the order in which this stop occurs in a typical
	// route trip, for comparisons with other stops matching the
	// same agency / route / stop / direction / headsign
	StopSequence int `json:"stop_sequence" db:"stop_sequence"`

	Dist      float64      `json:"dist" db:"dist" upsert:"omit"`
	Scheduled []*Departure `json:"scheduled" db:"-" upsert:"omit"`
	Live      []*Departure `json:"live" db:"-" upsert:"omit"`
}

// Table implements the upsert.Upserter interface, returning the table
// where we save stops.
func (s *Stop) Table() string {
	return "stop"
}

// Save saves a stop to the database
func (s *Stop) Save() error {
	_, err := upsert.Upsert(etc.DBConn, s)
	return err
}

// appendLive calls either the bus time API or the subway datamine API
// to add live info to our stop info.
func (s *Stop) appendLive(now time.Time) {
	route, err := GetRoute(s.AgencyID, s.RouteID, false)
	if err != nil {
		log.Println("can't load route", err)
		return
	}

	if route.Type == Bus {
		departures, err := GetLiveBus(
			s.RouteID, strconv.Itoa(s.DirectionID),
			s.ID,
		)
		if err != nil {
			log.Println("can't append live schedules")
			return
		}

		sort.Sort(departures)

		for i := 0; i < len(departures) && i < maxDepartures; i++ {
			s.Live = append(s.Live, departures[i])
		}
	} else if route.Type == Subway {

		departures, err := GetLiveSubways(
			s.RouteID, strconv.Itoa(s.DirectionID),
			s.ID,
		)
		if err != nil {
			log.Println("can't append live subway sched", err)
			return
		}

		sort.Sort(departures)
		for i := 0; i < len(departures) && i < maxDepartures; i++ {
			s.Live = append(s.Live, departures[i])
		}

	}
}

// String returns a descriptive string for this stop.
func (s Stop) String() string {
	return fmt.Sprintf("{%v %v %v %v %v @ (%v,%v)}",
		s.ID, s.Name, s.RouteID, s.Headsign, s.DirectionID, s.Lat, s.Lon,
	)
}

// Key() returns the unique string for this stop, so we can identify
// unique stops in the loader.
func (s Stop) Key() string {
	return fmt.Sprintf("%v%v", s.ID, s.RouteID)
}

// setDepartures checks the database and any relevant APIs to set the scheduled
// and live departures for this stop
func (s *Stop) setDepartures(now time.Time, db sqlx.Ext) (err error) {

	allDepartures := Departures{}

	yesterday := baseTime(now.Add(-time.Hour * 12))
	today := baseTime(now)

	yesterdayName := strings.ToLower(yesterday.Format("Monday"))
	todayName := strings.ToLower(now.Format("Monday"))

	func() {
		if yesterdayName != todayName {
			var yesterdayID string
			// Looks for trips starting yesterday that arrive here
			// after midnight
			yesterdayID, err = getServiceIDByDay(
				db, s.RouteID, yesterdayName, yesterday,
			)
			if err == sql.ErrNoRows {
				err = nil
				log.Println("no rows, ok, moving on")
				return
			}
			if err != nil {
				log.Println("can't get yesterday id", err)
				return
			}

			nowSecs :=
				now.Hour()*3600 + now.Minute()*60 + now.Second() + midnightSecs

			departures, err := getDepartures(
				s.AgencyID, s.RouteID, s.ID, yesterdayID,
				nowSecs, yesterday)
			if err != nil {
				log.Println("can't get departures", err)
				return
			}

			allDepartures = append(allDepartures, departures...)
		}
	}()

	func() {
		var todayID string
		todayID, err = getServiceIDByDay(db, s.RouteID, todayName, today)
		if err == sql.ErrNoRows {
			err = nil
			log.Println("no rows there", err)
			return
		}
		if err != nil {
			log.Println("can't get today id", err)
			return
		}

		nowSecs := now.Hour()*3600 + now.Minute()*60 + now.Second()

		departures, err := getDepartures(
			s.AgencyID, s.RouteID, s.ID, todayID,
			nowSecs, today)
		if err != nil {
			log.Println("can't get departures", err)
			return
		}

		allDepartures = append(allDepartures, departures...)

	}()

	sort.Sort(allDepartures)

	for i, d := range allDepartures {
		if i > maxDepartures {
			break
		}
		s.Scheduled = append(s.Scheduled, d)
	}

	// After reading scheduled times in the db, try to also append any
	// live info available
	s.appendLive(now)

	return
}

// GetStopsByQuery returns stops matching this StopQuery
func GetStopsByQuery(db sqlx.Ext, sq StopQuery) (stops []*Stop, err error) {
	now := time.Now()

	// Get rows matching the stop query
	rows, err := sqlx.NamedQuery(db, sq.Query(), sq)
	if err != nil {
		log.Println("can't get stops", sq.Query(), err)
		return
	}

	for rows.Next() {
		var stop Stop

		err = rows.StructScan(&stop)
		if err != nil {
			log.Println("can't scan stop", err)
			return
		}

		if sq.Departures {
			err = stop.setDepartures(now, db)
			if err != nil {
				log.Println("can't set departures", err)
				return
			}
		}

		stops = append(stops, &stop)
	}

	return
}

func getServiceIDByDay(db sqlx.Ext, routeID, day string, now time.Time) (serviceID string, err error) {

	// Select the service_id that:
	//   * matches our routeID and day
	//   * has an end_date after now
	//   * has a start_date before now
	//   * if there's more than one, choose the one with the latest start_date

	row := db.QueryRowx(`
		SELECT service_id
		FROM   service_route_day
		WHERE  day = $1 AND
			   end_date >= $2 AND
			   start_date <= $3 AND
			   route_id = $4
		ORDER BY start_date DESC
		LIMIT 1
	`,
		day, now, now, routeID,
	)

	err = row.Scan(&serviceID)
	if err != nil {
		log.Println("can't scan service id", err, day, now, routeID)
		return
	}

	return
}
<commit_msg>omit stop sequence from upsert<commit_after>package models

import (
	"database/sql"
	"fmt"
	"log"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/upsert"
	"github.com/jmoiron/sqlx"
)

// Stop is a single transit stop for a particular route. If a
// stop serves more than one route, there are multiple distinct
// entries for that stop.
type Stop struct {
	ID       string `json:"stop_id" db:"stop_id" upsert:"key"`
	RouteID  string `json:"route_id" db:"route_id" upsert:"key"`
	AgencyID string `json:"agency_id" db:"agency_id" upsert:"key"`
	Name     string `json:"stop_name" db:"stop_name"`

	DirectionID int    `json:"direction_id" db:"direction_id"`
	Headsign    string `json:"headsign" db:"headsign"`

	Lat float64 `json:"lat" db:"lat" upsert:"omit"`
	Lon float64 `json:"lon" db:"lon" upsert:"omit"`

	// Location is an "earth" field value that combines lat and lon into
	// a single field.
	Location interface{} `json:"-" db:"location" upsert_value:"ll_to_earth(:lat, :lon)"`

	// StopSequence is the order in which this stop occurs in a typical
	// route trip, for comparisons with other stops matching the
	// same agency / route / stop / direction / headsign
	StopSequence int `json:"stop_sequence" db:"stop_sequence" upsert:"omit"`

	Dist      float64      `json:"dist" db:"dist" upsert:"omit"`
	Scheduled []*Departure `json:"scheduled" db:"-" upsert:"omit"`
	Live      []*Departure `json:"live" db:"-" upsert:"omit"`
}

// Table implements the upsert.Upserter interface, returning the table
// where we save stops.
func (s *Stop) Table() string {
	return "stop"
}

// Save saves a stop to the database
func (s *Stop) Save() error {
	_, err := upsert.Upsert(etc.DBConn, s)
	return err
}

// appendLive calls either the bus time API or the subway datamine API
// to add live info to our stop info.
func (s *Stop) appendLive(now time.Time) {
	route, err := GetRoute(s.AgencyID, s.RouteID, false)
	if err != nil {
		log.Println("can't load route", err)
		return
	}

	if route.Type == Bus {
		departures, err := GetLiveBus(
			s.RouteID, strconv.Itoa(s.DirectionID),
			s.ID,
		)
		if err != nil {
			log.Println("can't append live schedules")
			return
		}

		sort.Sort(departures)

		for i := 0; i < len(departures) && i < maxDepartures; i++ {
			s.Live = append(s.Live, departures[i])
		}
	} else if route.Type == Subway {

		departures, err := GetLiveSubways(
			s.RouteID, strconv.Itoa(s.DirectionID),
			s.ID,
		)
		if err != nil {
			log.Println("can't append live subway sched", err)
			return
		}

		sort.Sort(departures)
		for i := 0; i < len(departures) && i < maxDepartures; i++ {
			s.Live = append(s.Live, departures[i])
		}

	}
}

// String returns a descriptive string for this stop.
func (s Stop) String() string {
	return fmt.Sprintf("{%v %v %v %v %v @ (%v,%v)}",
		s.ID, s.Name, s.RouteID, s.Headsign, s.DirectionID, s.Lat, s.Lon,
	)
}

// Key() returns the unique string for this stop, so we can identify
// unique stops in the loader.
func (s Stop) Key() string {
	return fmt.Sprintf("%v%v", s.ID, s.RouteID)
}

// setDepartures checks the database and any relevant APIs to set the scheduled
// and live departures for this stop
func (s *Stop) setDepartures(now time.Time, db sqlx.Ext) (err error) {

	allDepartures := Departures{}

	yesterday := baseTime(now.Add(-time.Hour * 12))
	today := baseTime(now)

	yesterdayName := strings.ToLower(yesterday.Format("Monday"))
	todayName := strings.ToLower(now.Format("Monday"))

	func() {
		if yesterdayName != todayName {
			var yesterdayID string
			// Looks for trips starting yesterday that arrive here
			// after midnight
			yesterdayID, err = getServiceIDByDay(
				db, s.RouteID, yesterdayName, yesterday,
			)
			if err == sql.ErrNoRows {
				err = nil
				log.Println("no rows, ok, moving on")
				return
			}
			if err != nil {
				log.Println("can't get yesterday id", err)
				return
			}

			nowSecs :=
				now.Hour()*3600 + now.Minute()*60 + now.Second() + midnightSecs

			departures, err := getDepartures(
				s.AgencyID, s.RouteID, s.ID, yesterdayID,
				nowSecs, yesterday)
			if err != nil {
				log.Println("can't get departures", err)
				return
			}

			allDepartures = append(allDepartures, departures...)
		}
	}()

	func() {
		var todayID string
		todayID, err = getServiceIDByDay(db, s.RouteID, todayName, today)
		if err == sql.ErrNoRows {
			err = nil
			log.Println("no rows there", err)
			return
		}
		if err != nil {
			log.Println("can't get today id", err)
			return
		}

		nowSecs := now.Hour()*3600 + now.Minute()*60 + now.Second()

		departures, err := getDepartures(
			s.AgencyID, s.RouteID, s.ID, todayID,
			nowSecs, today)
		if err != nil {
			log.Println("can't get departures", err)
			return
		}

		allDepartures = append(allDepartures, departures...)

	}()

	sort.Sort(allDepartures)

	for i, d := range allDepartures {
		if i > maxDepartures {
			break
		}
		s.Scheduled = append(s.Scheduled, d)
	}

	// After reading scheduled times in the db, try to also append any
	// live info available
	s.appendLive(now)

	return
}

// GetStopsByQuery returns stops matching this StopQuery
func GetStopsByQuery(db sqlx.Ext, sq StopQuery) (stops []*Stop, err error) {
	now := time.Now()

	// Get rows matching the stop query
	rows, err := sqlx.NamedQuery(db, sq.Query(), sq)
	if err != nil {
		log.Println("can't get stops", sq.Query(), err)
		return
	}

	for rows.Next() {
		var stop Stop

		err = rows.StructScan(&stop)
		if err != nil {
			log.Println("can't scan stop", err)
			return
		}

		if sq.Departures {
			err = stop.setDepartures(now, db)
			if err != nil {
				log.Println("can't set departures", err)
				return
			}
		}

		stops = append(stops, &stop)
	}

	return
}

func getServiceIDByDay(db sqlx.Ext, routeID, day string, now time.Time) (serviceID string, err error) {

	// Select the service_id that:
	//   * matches our routeID and day
	//   * has an end_date after now
	//   * has a start_date before now
	//   * if there's more than one, choose the one with the latest start_date

	row := db.QueryRowx(`
		SELECT service_id
		FROM   service_route_day
		WHERE  day = $1 AND
			   end_date >= $2 AND
			   start_date <= $3 AND
			   route_id = $4
		ORDER BY start_date DESC
		LIMIT 1
	`,
		day, now, now, routeID,
	)

	err = row.Scan(&serviceID)
	if err != nil {
		log.Println("can't scan service id", err, day, now, routeID)
		return
	}

	return
}
<|endoftext|>