<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

var (
	twitterUsername       string
	twitterConsumerKey    string
	twitterConsumerSecret string
	twitterAuthToken      string
	twitterAuthSecret     string

	twitterMentionRegex = regexp.MustCompile("^@\\w+\\s*")
	twitterTextRegex    = regexp.MustCompile("@\\w+|\\s+|.?")
	twitterAPIClient    *twitter.Client
	twitterUploadClient *http.Client
)

const (
	maxTweetLen              = 140
	groupThreshold           = 0.8
	twitterUploadURL         = "https://upload.twitter.com/1.1/media/upload.json"
	twitterUploadMetadataURL = "https://upload.twitter.com/1.1/media/metadata/create.json"
)

type twitterPlugin struct{}

func (p twitterPlugin) EnvVariables() []EnvVariable {
	return []EnvVariable{
		{
			Name:     "TWITTER_USERNAME",
			Variable: &twitterUsername,
		},
		{
			Name:     "TWITTER_CONSUMER_KEY",
			Variable: &twitterConsumerKey,
		},
		{
			Name:     "TWITTER_CONSUMER_SECRET",
			Variable: &twitterConsumerSecret,
		},
		{
			Name:     "TWITTER_ACCESS_TOKEN",
			Variable: &twitterAuthToken,
		},
		{
			Name:     "TWITTER_ACCESS_TOKEN_SECRET",
			Variable: &twitterAuthSecret,
		},
	}
}

func (p twitterPlugin) Name() string {
	return "twitter"
}

func NewTwitterPlugin() WorkerPlugin {
	return twitterPlugin{}
}

func (p twitterPlugin) Start(ch chan error) {
	defer close(ch)

	config := oauth1.NewConfig(twitterConsumerKey, twitterConsumerSecret)
	token := oauth1.NewToken(twitterAuthToken, twitterAuthSecret)

	httpClient := config.Client(oauth1.NoContext, token)
	twitterUploadClient = httpClient
	twitterAPIClient = twitter.NewClient(httpClient)

	params := &twitter.StreamUserParams{
		With:          "user",
		StallWarnings: twitter.Bool(true),
	}
	stream, err := twitterAPIClient.Streams.User(params)
	if err != nil {
		ch <- err
		return
	}

	demux := twitter.NewSwitchDemux()
	demux.Tweet = func(tweet *twitter.Tweet) {
		handleTweet(tweet, ch)
	}
	demux.DM = handleDM
	demux.StreamLimit = handleStreamLimit
	demux.StreamDisconnect = handleStreamDisconnect
	demux.Warning = handleWarning
	demux.Other = handleOther

	demux.HandleChan(stream.Messages)
}

func logMessage(msg interface{}, desc string) {
	if msgJSON, err := json.MarshalIndent(msg, "", "  "); err == nil {
		log.Printf("Received %s: %s\n", desc, string(msgJSON[:]))
	} else {
		logMessageStruct(msg, desc)
	}
}

func logMessageStruct(msg interface{}, desc string) {
	log.Printf("Received %s: %+v\n", desc, msg)
}

func trimReply(t string) string {
	for twitterMentionRegex.MatchString(t) {
		t = twitterMentionRegex.ReplaceAllString(t, "")
	}
	return t
}

func transformTwitterText(t string) string {
	var buffer bytes.Buffer
	letters := twitterTextRegex.FindAllString(t, -1)
	trFuncs := []func(string) string{
		strings.ToUpper,
		strings.ToLower,
	}
	idx := rand.Intn(2)
	groupSize := rand.Intn(2) + 1
	for _, ch := range letters {
		// ignore twitter usernames
		if len(ch) == 1 && strings.TrimSpace(ch) != "" {
			ch = trFuncs[idx](ch)
			groupSize--
			if groupSize == 0 {
				idx = (idx + 1) % 2
				groupSize = 1
				if rand.Float64() > groupThreshold {
					groupSize++
				}
			}
		}
		buffer.WriteString(ch)
	}

	return buffer.String()
}

type twitterImageData struct {
	ImageType string `json:"image_type"`
	Width     int    `json:"w"`
	Height    int    `json:"h"`
}

type twitterUploadResponse struct {
	MediaID          int64             `json:"media_id"`
	MediaIDStr       string            `json:"media_id_string"`
	Size             int               `json:"size"`
	ExpiresAfterSecs int               `json:"expires_after_secs"`
	Image            *twitterImageData `json:"image"`
}

func uploadImage() (int64, string, error) {
	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	memeFile, err := os.Open(memePath)
	if err != nil {
		return 0, "", fmt.Errorf("opening meme image file error: %s", err)
	}
	defer memeFile.Close()

	fw, err := w.CreateFormFile("media", filepath.Base(memePath))
	if err != nil {
		return 0, "", fmt.Errorf("creating multipart form file header error: %s", err)
	}
	if _, err = io.Copy(fw, memeFile); err != nil {
		return 0, "", fmt.Errorf("io copy error: %s", err)
	}
	w.Close()

	req, err := http.NewRequest("POST", twitterUploadURL, &b)
	if err != nil {
		return 0, "", fmt.Errorf("creating POST request error: %s", err)
	}
	req.Header.Set("Content-Type", w.FormDataContentType())

	res, err := twitterUploadClient.Do(req)
	if err != nil {
		return 0, "", fmt.Errorf("sending POST request error: %s", err)
	}

	id, idStr, err := parseUploadResponse(res)
	if err != nil {
		return 0, "", err
	}

	return id, idStr, nil
}

func parseUploadResponse(res *http.Response) (int64, string, error) {
	if res.StatusCode != http.StatusOK {
		return 0, "", fmt.Errorf("image upload bad status: %s", res.Status)
	}
	defer res.Body.Close()

	var resBuf bytes.Buffer
	if _, err := resBuf.ReadFrom(res.Body); err != nil {
		return 0, "", fmt.Errorf("reading from http response body error: %s", err)
	}
	log.Println("http body:", resBuf.String())

	resp := twitterUploadResponse{}
	if err := json.Unmarshal(resBuf.Bytes(), &resp); err != nil {
		return 0, "", fmt.Errorf("unmarshalling twitter upload response error: %s", err)
	}
	log.Printf("json struct: %+v\n", resp)

	// TODO: add logic dealing with the expires_after_secs
	return resp.MediaID, resp.MediaIDStr, nil
}

type twitterAltText struct {
	Text string `json:"text"`
}

type twitterImageMetadata struct {
	MediaID string          `json:"media_id"`
	AltText *twitterAltText `json:"alt_text"`
}

func uploadMetadata(mediaID, text string) error {
	md := twitterImageMetadata{
		MediaID: mediaID,
		AltText: &twitterAltText{
			Text: text,
		},
	}
	raw, err := json.Marshal(md)
	if err != nil {
		return fmt.Errorf("json marshal error: %s", err)
	}
	req, err := http.NewRequest("POST", twitterUploadMetadataURL, bytes.NewReader(raw))
	if err != nil {
		return fmt.Errorf("making http request error: %s", err)
	}
	req.Header.Set("Content-Type", "application/json; charset=UTF-8")

	res, err := twitterUploadClient.Do(req)
	if err != nil {
		return fmt.Errorf("sending POST request error: %s", err)
	}

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("metadata upload returned status code %d", res.StatusCode)
	}

	return nil
}

func handleTweet(tweet *twitter.Tweet, ch chan error) {
	logMessageStruct(tweet, "Tweet")

	var tt string
	if tweet.InReplyToStatusIDStr == "" {
		// case where someone tweets @ the bot

		tt = tweet.Text
	} else if tweet.User.ScreenName != twitterUsername {
		params := twitter.StatusLookupParams{
			TrimUser:        twitter.Bool(true),
			IncludeEntities: twitter.Bool(false),
		}
		tweets, resp, err := twitterAPIClient.Statuses.Lookup([]int64{tweet.InReplyToStatusID}, &params)
		if err != nil {
			ch <- fmt.Errorf("status lookup error: %s", err)
			return
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			ch <- fmt.Errorf("status lookup HTTP status code: %d", resp.StatusCode)
			return
		}
		if len(tweets) == 0 {
			ch <- errors.New("number of returned tweets is 0")
			return
		}
		tt = tweets[0].Text
	} else {
		return
	}

	rt := fmt.Sprintf("@%s %s", tweet.User.ScreenName, transformTwitterText(tt))
	if len(rt) > maxTweetLen {
		log.Println("Exceeded max tweet length:", len(rt), rt)
		rt = fmt.Sprintf("@%s %s", tweet.User.ScreenName, transformTwitterText(trimReply(tt)))
	}
	mediaID, mediaIDStr, err := uploadImage()
	if err != nil {
		ch <- fmt.Errorf("upload image error: %s", err)
		return
	}
	if err = uploadMetadata(mediaIDStr, tt); err != nil {
		// we can continue from a metadata upload error
		// because it is not essential
		ch <- fmt.Errorf("metadata upload error: %s", err)
	}

	params := twitter.StatusUpdateParams{
		InReplyToStatusID: tweet.ID,
		TrimUser:          twitter.Bool(true),
		MediaIds:          []int64{mediaID},
	}
	_, resp, err := twitterAPIClient.Statuses.Update(rt, &params)
	if err != nil {
		ch <- fmt.Errorf("status update error: %s", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		ch <- fmt.Errorf("response tweet status code: %d", resp.StatusCode)
		return
	}
}

func handleDM(dm *twitter.DirectMessage) {
	logMessage(dm, "DM")
}

func handleStreamLimit(sl *twitter.StreamLimit) {
	logMessage(sl, "stream limit message")
}

func handleStreamDisconnect(sd *twitter.StreamDisconnect) {
	logMessage(sd, "stream disconnect message")
}

func handleWarning(w *twitter.StallWarning) {
	logMessage(w, "stall warning")
}

func handleOther(message interface{}) {
	logMessage(message, `"other" message type`)
}
<commit_msg>Update reply policy<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

var (
	twitterUsername       string
	twitterConsumerKey    string
	twitterConsumerSecret string
	twitterAuthToken      string
	twitterAuthSecret     string

	twitterMentionRegex = regexp.MustCompile("^@\\w+\\s*")
	twitterTextRegex    = regexp.MustCompile("@\\w+|\\s+|.?")
	twitterAPIClient    *twitter.Client
	twitterUploadClient *http.Client
)

const (
	maxTweetLen              = 140
	groupThreshold           = 0.8
	twitterUploadURL         = "https://upload.twitter.com/1.1/media/upload.json"
	twitterUploadMetadataURL = "https://upload.twitter.com/1.1/media/metadata/create.json"
)

type twitterPlugin struct{}

func (p twitterPlugin) EnvVariables() []EnvVariable {
	return []EnvVariable{
		{
			Name:     "TWITTER_USERNAME",
			Variable: &twitterUsername,
		},
		{
			Name:     "TWITTER_CONSUMER_KEY",
			Variable: &twitterConsumerKey,
		},
		{
			Name:     "TWITTER_CONSUMER_SECRET",
			Variable: &twitterConsumerSecret,
		},
		{
			Name:     "TWITTER_ACCESS_TOKEN",
			Variable: &twitterAuthToken,
		},
		{
			Name:     "TWITTER_ACCESS_TOKEN_SECRET",
			Variable: &twitterAuthSecret,
		},
	}
}

func (p twitterPlugin) Name() string {
	return "twitter"
}

func NewTwitterPlugin() WorkerPlugin {
	return twitterPlugin{}
}

func (p twitterPlugin) Start(ch chan error) {
	defer close(ch)

	config := oauth1.NewConfig(twitterConsumerKey, twitterConsumerSecret)
	token := oauth1.NewToken(twitterAuthToken, twitterAuthSecret)

	httpClient := config.Client(oauth1.NoContext, token)
	twitterUploadClient = httpClient
	twitterAPIClient = twitter.NewClient(httpClient)

	params := &twitter.StreamUserParams{
		With:          "user",
		StallWarnings: twitter.Bool(true),
	}
	stream, err := twitterAPIClient.Streams.User(params)
	if err != nil {
		ch <- err
		return
	}

	demux := twitter.NewSwitchDemux()
	demux.Tweet = func(tweet *twitter.Tweet) {
		handleTweet(tweet, ch)
	}
	demux.DM = handleDM
	demux.StreamLimit = handleStreamLimit
	demux.StreamDisconnect = handleStreamDisconnect
	demux.Warning = handleWarning
	demux.Other = handleOther

	demux.HandleChan(stream.Messages)
}

func logMessage(msg interface{}, desc string) {
	if msgJSON, err := json.MarshalIndent(msg, "", "  "); err == nil {
		log.Printf("Received %s: %s\n", desc, string(msgJSON[:]))
	} else {
		logMessageStruct(msg, desc)
	}
}

func logMessageStruct(msg interface{}, desc string) {
	log.Printf("Received %s: %+v\n", desc, msg)
}

func trimReply(t string) string {
	for twitterMentionRegex.MatchString(t) {
		t = twitterMentionRegex.ReplaceAllString(t, "")
	}
	return t
}

func transformTwitterText(t string) string {
	var buffer bytes.Buffer
	letters := twitterTextRegex.FindAllString(t, -1)
	trFuncs := []func(string) string{
		strings.ToUpper,
		strings.ToLower,
	}
	idx := rand.Intn(2)
	groupSize := rand.Intn(2) + 1
	for _, ch := range letters {
		// ignore twitter usernames
		if len(ch) == 1 && strings.TrimSpace(ch) != "" {
			ch = trFuncs[idx](ch)
			groupSize--
			if groupSize == 0 {
				idx = (idx + 1) % 2
				groupSize = 1
				if rand.Float64() > groupThreshold {
					groupSize++
				}
			}
		}
		buffer.WriteString(ch)
	}

	return buffer.String()
}

type twitterImageData struct {
	ImageType string `json:"image_type"`
	Width     int    `json:"w"`
	Height    int    `json:"h"`
}

type twitterUploadResponse struct {
	MediaID          int64             `json:"media_id"`
	MediaIDStr       string            `json:"media_id_string"`
	Size             int               `json:"size"`
	ExpiresAfterSecs int               `json:"expires_after_secs"`
	Image            *twitterImageData `json:"image"`
}

func uploadImage() (int64, string, error) {
	var b bytes.Buffer
	w := multipart.NewWriter(&b)
	memeFile, err := os.Open(memePath)
	if err != nil {
		return 0, "", fmt.Errorf("opening meme image file error: %s", err)
	}
	defer memeFile.Close()

	fw, err := w.CreateFormFile("media", filepath.Base(memePath))
	if err != nil {
		return 0, "", fmt.Errorf("creating multipart form file header error: %s", err)
	}
	if _, err = io.Copy(fw, memeFile); err != nil {
		return 0, "", fmt.Errorf("io copy error: %s", err)
	}
	w.Close()

	req, err := http.NewRequest("POST", twitterUploadURL, &b)
	if err != nil {
		return 0, "", fmt.Errorf("creating POST request error: %s", err)
	}
	req.Header.Set("Content-Type", w.FormDataContentType())

	res, err := twitterUploadClient.Do(req)
	if err != nil {
		return 0, "", fmt.Errorf("sending POST request error: %s", err)
	}

	id, idStr, err := parseUploadResponse(res)
	if err != nil {
		return 0, "", err
	}

	return id, idStr, nil
}

func parseUploadResponse(res *http.Response) (int64, string, error) {
	if res.StatusCode != http.StatusOK {
		return 0, "", fmt.Errorf("image upload bad status: %s", res.Status)
	}
	defer res.Body.Close()

	var resBuf bytes.Buffer
	if _, err := resBuf.ReadFrom(res.Body); err != nil {
		return 0, "", fmt.Errorf("reading from http response body error: %s", err)
	}
	log.Println("http body:", resBuf.String())

	resp := twitterUploadResponse{}
	if err := json.Unmarshal(resBuf.Bytes(), &resp); err != nil {
		return 0, "", fmt.Errorf("unmarshalling twitter upload response error: %s", err)
	}
	log.Printf("json struct: %+v\n", resp)

	// TODO: add logic dealing with the expires_after_secs
	return resp.MediaID, resp.MediaIDStr, nil
}

type twitterAltText struct {
	Text string `json:"text"`
}

type twitterImageMetadata struct {
	MediaID string          `json:"media_id"`
	AltText *twitterAltText `json:"alt_text"`
}

func uploadMetadata(mediaID, text string) error {
	md := twitterImageMetadata{
		MediaID: mediaID,
		AltText: &twitterAltText{
			Text: text,
		},
	}
	raw, err := json.Marshal(md)
	if err != nil {
		return fmt.Errorf("json marshal error: %s", err)
	}
	req, err := http.NewRequest("POST", twitterUploadMetadataURL, bytes.NewReader(raw))
	if err != nil {
		return fmt.Errorf("making http request error: %s", err)
	}
	req.Header.Set("Content-Type", "application/json; charset=UTF-8")

	res, err := twitterUploadClient.Do(req)
	if err != nil {
		return fmt.Errorf("sending POST request error: %s", err)
	}

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("metadata upload returned status code %d", res.StatusCode)
	}

	return nil
}

func handleTweet(tweet *twitter.Tweet, ch chan error) {
	logMessageStruct(tweet, "Tweet")

	var tt string
	if tweet.InReplyToStatusIDStr == "" {
		// case where someone tweets @ the bot

		tt = tweet.Text
	} else if tweet.User.ScreenName != twitterUsername {
		params := twitter.StatusLookupParams{
			IncludeEntities: twitter.Bool(false),
		}
		tweets, resp, err := twitterAPIClient.Statuses.Lookup([]int64{tweet.InReplyToStatusID}, &params)
		if err != nil {
			ch <- fmt.Errorf("status lookup error: %s", err)
			return
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			ch <- fmt.Errorf("status lookup HTTP status code: %d", resp.StatusCode)
			return
		}
		if len(tweets) == 0 {
			ch <- errors.New("number of returned tweets is 0")
			return
		}
		tt = fmt.Sprintf("@%s %s", tweets[0].User.ScreenName, trimReply(tweets[0].Text))
	} else {
		return
	}

	rt := fmt.Sprintf("@%s %s", tweet.User.ScreenName, transformTwitterText(tt))
	if len(rt) > maxTweetLen {
		log.Println("Exceeded max tweet length:", len(rt), rt)
		rt = fmt.Sprintf("@%s %s", tweet.User.ScreenName, transformTwitterText(trimReply(tt)))
	}
	mediaID, mediaIDStr, err := uploadImage()
	if err != nil {
		ch <- fmt.Errorf("upload image error: %s", err)
		return
	}
	if err = uploadMetadata(mediaIDStr, tt); err != nil {
		// we can continue from a metadata upload error
		// because it is not essential
		ch <- fmt.Errorf("metadata upload error: %s", err)
	}

	params := twitter.StatusUpdateParams{
		InReplyToStatusID: tweet.ID,
		TrimUser:          twitter.Bool(true),
		MediaIds:          []int64{mediaID},
	}
	_, resp, err := twitterAPIClient.Statuses.Update(rt, &params)
	if err != nil {
		ch <- fmt.Errorf("status update error: %s", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		ch <- fmt.Errorf("response tweet status code: %d", resp.StatusCode)
		return
	}
}

func handleDM(dm *twitter.DirectMessage) {
	logMessage(dm, "DM")
}

func handleStreamLimit(sl *twitter.StreamLimit) {
	logMessage(sl, "stream limit message")
}

func handleStreamDisconnect(sd *twitter.StreamDisconnect) {
	logMessage(sd, "stream disconnect message")
}

func handleWarning(w *twitter.StallWarning) {
	logMessage(w, "stall warning")
}

func handleOther(message interface{}) {
	logMessage(message, `"other" message type`)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"strings"

	"github.com/brnstz/bus/common"
	"github.com/brnstz/bus/loader"
	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
)

func doOne(dir string, db *sqlx.DB) {

	stype := ""
	if strings.HasSuffix(dir, "subway/") {
		stype = "subway"
	} else {
		stype = "bus"
	}

	log.Println(dir, stype)
	l := loader.NewLoader(dir)

	for _, s := range l.ServiceRouteDays {
		log.Println("Inserting service route day: ", s)
		_, err := db.Exec(`
				INSERT INTO service_route_day
				(route_id, service_id, day, start_date, end_date)
				VALUES($1, $2, $3, $4, $5)
			`, s.RouteId, s.ServiceId, s.Day, s.StartDate, s.EndDate,
		)
		if err != nil {
			log.Fatal(err)
		}
	}

	for _, s := range l.Stops {
		log.Println("Inserting stop: ", s)
		_, err := db.Exec(`
				INSERT INTO stop
				(stop_id, stop_name, direction_id, headsign, route_id,
				 location, stype)
				VALUES($1, $2, $3, $4, $5, ll_to_earth($6, $7), $8)
			`,
			s.Id, s.Name, s.DirectionId, s.Headsign, s.RouteId,
			s.Lat, s.Lon, stype,
		)

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

	for _, s := range l.ScheduledStopTimes {
		log.Println("Inserting scheduled stop time: ", s)
		_, err := db.Exec(`
				INSERT INTO scheduled_stop_time
				(route_id, stop_id, service_id, departure_sec)
				VALUES($1, $2, $3, $4)
			`, s.RouteId, s.StopId, s.ServiceId, s.DepartureSec,
		)
		if err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	db := common.DB

	for _, dir := range []string{
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/subway/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/brooklyn/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/manhattan/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/queens/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/staten_island/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/bronx/",
	} {

		doOne(dir, db)
	}

	log.Println("finished all boroughs")
}
<commit_msg>err handling and stack trace<commit_after>package main

import (
	"log"
	"os"
	"os/signal"
	"strings"

	"github.com/brnstz/bus/common"
	"github.com/brnstz/bus/loader"
	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
)

func doOne(dir string, db *sqlx.DB) {

	stype := ""
	if strings.HasSuffix(dir, "subway/") {
		stype = "subway"
	} else {
		stype = "bus"
	}

	log.Println(dir, stype)
	l := loader.NewLoader(dir)

	for _, s := range l.ServiceRouteDays {
		log.Println("Inserting service route day: ", s)
		_, err := db.Exec(`
				INSERT INTO service_route_day
				(route_id, service_id, day, start_date, end_date)
				VALUES($1, $2, $3, $4, $5)
			`, s.RouteId, s.ServiceId, s.Day, s.StartDate, s.EndDate,
		)
		if err != nil {
			log.Println("ERROR SERVICE ROUTE DAYS: ", err, s)
		}
	}

	for _, s := range l.Stops {
		log.Println("Inserting stop: ", s)
		_, err := db.Exec(`
				INSERT INTO stop
				(stop_id, stop_name, direction_id, headsign, route_id,
				 location, stype)
				VALUES($1, $2, $3, $4, $5, ll_to_earth($6, $7), $8)
			`,
			s.Id, s.Name, s.DirectionId, s.Headsign, s.RouteId,
			s.Lat, s.Lon, stype,
		)

		if err != nil {
			log.Println("ERROR STOPS: ", err, s)
		}
	}

	for _, s := range l.ScheduledStopTimes {
		log.Println("Inserting scheduled stop time: ", s)
		_, err := db.Exec(`
				INSERT INTO scheduled_stop_time
				(route_id, stop_id, service_id, departure_sec)
				VALUES($1, $2, $3, $4)
			`, s.RouteId, s.StopId, s.ServiceId, s.DepartureSec,
		)
		if err != nil {
			log.Println("ERROR SCHEDULED STOP TIMES: ", err, s)
		}
	}
}

func main() {
	db := common.DB

	// Dump stack trace on kill
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for sig := range c {
			log.Println("got signal", sig)
			panic("stack trace here")
		}
	}()

	for _, dir := range []string{
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/subway/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/brooklyn/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/manhattan/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/queens/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/staten_island/",
		"/Users/bseitz/go/src/github.com/brnstz/bus/schema/bronx/",
	} {

		doOne(dir, db)
	}

	log.Println("finished all boroughs")
}
<|endoftext|>


<commit_before>package operation

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"path/filepath"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/security"
)

type UploadResult struct {
	Name  string `json:"name,omitempty"`
	Size  uint32 `json:"size,omitempty"`
	Error string `json:"error,omitempty"`
	ETag  string `json:"eTag,omitempty"`
}

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}}
}

var fileNameEscaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")

// Upload sends a POST request to a volume server to upload the content
func Upload(uploadUrl string, filename string, reader io.Reader, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
	return upload_content(uploadUrl, func(w io.Writer) (err error) {
		_, err = io.Copy(w, reader)
		return
	}, filename, isGzipped, mtype, pairMap, jwt)
}
func upload_content(uploadUrl string, fillBufferFunction func(w io.Writer) error, filename string, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
	body_buf := bytes.NewBufferString("")
	body_writer := multipart.NewWriter(body_buf)
	h := make(textproto.MIMEHeader)
	h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, fileNameEscaper.Replace(filename)))
	if mtype == "" {
		mtype = mime.TypeByExtension(strings.ToLower(filepath.Ext(filename)))
	}
	if mtype != "" {
		h.Set("Content-Type", mtype)
	}
	if isGzipped {
		h.Set("Content-Encoding", "gzip")
	}

	file_writer, cp_err := body_writer.CreatePart(h)
	if cp_err != nil {
		glog.V(0).Infoln("error creating form file", cp_err.Error())
		return nil, cp_err
	}
	if err := fillBufferFunction(file_writer); err != nil {
		glog.V(0).Infoln("error copying data", err)
		return nil, err
	}
	content_type := body_writer.FormDataContentType()
	if err := body_writer.Close(); err != nil {
		glog.V(0).Infoln("error closing body", err)
		return nil, err
	}

	req, postErr := http.NewRequest("POST", uploadUrl, body_buf)
	if postErr != nil {
		glog.V(0).Infoln("failing to upload to", uploadUrl, postErr.Error())
		return nil, postErr
	}
	req.Header.Set("Content-Type", content_type)
	for k, v := range pairMap {
		req.Header.Set(k, v)
	}
	if jwt != "" {
		req.Header.Set("Authorization", "BEARER "+string(jwt))
	}
	resp, post_err := client.Do(req)
	if post_err != nil {
		glog.V(0).Infoln("failing to upload to", uploadUrl, post_err.Error())
		return nil, post_err
	}
	defer resp.Body.Close()

	if resp.StatusCode < http.StatusOK ||
		resp.StatusCode > http.StatusIMUsed {
		return nil, errors.New(http.StatusText(resp.StatusCode))
	}

	etag := getEtag(resp)
	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		return nil, ra_err
	}
	var ret UploadResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.V(0).Infoln("failing to read upload response", uploadUrl, string(resp_body))
		return nil, unmarshal_err
	}
	if ret.Error != "" {
		return nil, errors.New(ret.Error)
	}
	ret.ETag = etag
	return &ret, nil
}

func getEtag(r *http.Response) (etag string) {
	etag = r.Header.Get("ETag")
	if strings.HasPrefix(etag, "\"") && strings.HasSuffix(etag, "\"") {
		etag = etag[1 : len(etag)-1]
	}
	return
}
<commit_msg>Revert "fix https://github.com/chrislusf/seaweedfs/issues/780"<commit_after>package operation

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"path/filepath"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/security"
)

type UploadResult struct {
	Name  string `json:"name,omitempty"`
	Size  uint32 `json:"size,omitempty"`
	Error string `json:"error,omitempty"`
	ETag  string `json:"eTag,omitempty"`
}

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}}
}

var fileNameEscaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"")

// Upload sends a POST request to a volume server to upload the content
func Upload(uploadUrl string, filename string, reader io.Reader, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
	return upload_content(uploadUrl, func(w io.Writer) (err error) {
		_, err = io.Copy(w, reader)
		return
	}, filename, isGzipped, mtype, pairMap, jwt)
}
func upload_content(uploadUrl string, fillBufferFunction func(w io.Writer) error, filename string, isGzipped bool, mtype string, pairMap map[string]string, jwt security.EncodedJwt) (*UploadResult, error) {
	body_buf := bytes.NewBufferString("")
	body_writer := multipart.NewWriter(body_buf)
	h := make(textproto.MIMEHeader)
	h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, fileNameEscaper.Replace(filename)))
	if mtype == "" {
		mtype = mime.TypeByExtension(strings.ToLower(filepath.Ext(filename)))
	}
	if mtype != "" {
		h.Set("Content-Type", mtype)
	}
	if isGzipped {
		h.Set("Content-Encoding", "gzip")
	}

	file_writer, cp_err := body_writer.CreatePart(h)
	if cp_err != nil {
		glog.V(0).Infoln("error creating form file", cp_err.Error())
		return nil, cp_err
	}
	if err := fillBufferFunction(file_writer); err != nil {
		glog.V(0).Infoln("error copying data", err)
		return nil, err
	}
	content_type := body_writer.FormDataContentType()
	if err := body_writer.Close(); err != nil {
		glog.V(0).Infoln("error closing body", err)
		return nil, err
	}

	req, postErr := http.NewRequest("POST", uploadUrl, body_buf)
	if postErr != nil {
		glog.V(0).Infoln("failing to upload to", uploadUrl, postErr.Error())
		return nil, postErr
	}
	req.Header.Set("Content-Type", content_type)
	for k, v := range pairMap {
		req.Header.Set(k, v)
	}
	if jwt != "" {
		req.Header.Set("Authorization", "BEARER "+string(jwt))
	}
	resp, post_err := client.Do(req)
	if post_err != nil {
		glog.V(0).Infoln("failing to upload to", uploadUrl, post_err.Error())
		return nil, post_err
	}
	defer resp.Body.Close()
	etag := getEtag(resp)
	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		return nil, ra_err
	}
	var ret UploadResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.V(0).Infoln("failing to read upload response", uploadUrl, string(resp_body))
		return nil, unmarshal_err
	}
	if ret.Error != "" {
		return nil, errors.New(ret.Error)
	}
	ret.ETag = etag
	return &ret, nil
}

func getEtag(r *http.Response) (etag string) {
	etag = r.Header.Get("ETag")
	if strings.HasPrefix(etag, "\"") && strings.HasSuffix(etag, "\"") {
		etag = etag[1 : len(etag)-1]
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"log"
	"os"

	_ "github.com/denisenkom/go-mssqldb"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	"github.com/spf13/pflag"

	"github.com/go-testfixtures/testfixtures/v3"
)

var version = "master"

func main() {
	log.SetFlags(0)
	log.SetOutput(os.Stderr)

	var (
		versionFlag           bool
		dialect               string
		connString            string
		dir                   string
		files                 []string
		paths                 []string
		useAlterContraint     bool
		skipResetSequences    bool
		resetSequencesTo      int64
		skipTestDatabaseCheck bool
	)

	pflag.BoolVar(&versionFlag, "version", false, "show testfixtures version")
	pflag.StringVarP(&dialect, "dialect", "d", "", "which database system you're using (postgres, timescaledb, mysql, mariadb, sqlite or sqlserver)")
	pflag.StringVarP(&connString, "conn", "c", "", "a database connection string")
	pflag.StringVarP(&dir, "dir", "D", "", "a directory of YAML fixtures to load")
	pflag.StringSliceVarP(&files, "files", "f", nil, "a list of YAML files to load")
	pflag.StringSliceVarP(&paths, "paths", "p", nil, "a list of fixture paths to load (directory or file)")
	pflag.BoolVar(&useAlterContraint, "alter-constraint", false, "use ALTER CONSTRAINT to disable referential integrity (PostgreSQL only)")
	pflag.BoolVar(&skipResetSequences, "no-reset-sequences", false, "skip reset of sequences after loading (PostgreSQL only)")
	pflag.Int64Var(&resetSequencesTo, "reset-sequences-to", 0, "sets the number sequences will be reset after loading fixtures (PostgreSQL only, defaults to 10000)")
	pflag.BoolVar(&skipTestDatabaseCheck, "dangerous-no-test-database-check", false, `skips check for "test" in database name (use with caution)`)
	pflag.Parse()

	if versionFlag {
		log.Printf("testfixtures version: %s", version)
		return
	}

	if dialect == "" && connString == "" {
		log.Fatal("testfixtures: both --dialect (-d) and --conn (-c) are required")
		return
	}
	if dir == "" && len(files) == 0 && len(paths) == 0 {
		log.Fatal("testfixtures: either --dir (-D) or --files (-f) or --paths (-p) need to be given")
		return
	}

	dialect, err := getDialect(dialect)
	if err != nil {
		log.Fatal(err)
		return
	}

	db, err := sql.Open(dialect, connString)
	if err != nil {
		log.Fatalf("testfixtures: could not connect to database: %v", err)
		return
	}
	defer db.Close()

	if err := db.Ping(); err != nil {
		log.Fatalf("testfixtures: could not ping database: %v", err)
		return
	}

	options := []func(*testfixtures.Loader) error{
		testfixtures.Database(db),
		testfixtures.Dialect(dialect),
	}
	if dir != "" {
		options = append(options, testfixtures.Directory(dir))
	}
	if len(files) > 0 {
		options = append(options, testfixtures.Files(files...))
	}
	if len(paths) > 0 {
		options = append(options, testfixtures.Paths(paths...))
	}
	if useAlterContraint {
		options = append(options, testfixtures.UseAlterConstraint())
	}
	if skipResetSequences {
		options = append(options, testfixtures.SkipResetSequences())
	}
	if resetSequencesTo > 0 {
		options = append(options, testfixtures.ResetSequencesTo(resetSequencesTo))
	}
	if skipTestDatabaseCheck {
		options = append(options, testfixtures.DangerousSkipTestDatabaseCheck())
	}

	loader, err := testfixtures.New(options...)
	if err != nil {
		log.Fatal(err)
	}
	if err := loader.Load(); err != nil {
		log.Fatal(err)
	}
	log.Printf("testfixtures: fixtures loaded successfully")
}

func getDialect(dialect string) (string, error) {
	switch dialect {
	case "postgres", "postgresql", "timescaledb":
		return "postgres", nil
	case "mysql", "mariadb":
		return "mysql", nil
	case "sqlite", "sqlite3":
		if !isSQLiteSupported() {
			return "", fmt.Errorf("testfixtures: SQLite is not supported in this build")
		}
		return "sqlite3", nil
	case "mssql", "sqlserver":
		return "sqlserver", nil
	default:
		return "", fmt.Errorf(`testfixtures: unrecognized dialect "%s"`, dialect)
	}
}

func isSQLiteSupported() bool {
	for _, d := range sql.Drivers() {
		if d == "sqlite3" {
			return true
		}
	}
	return false
}
<commit_msg>[#63] testfixtures cli add dumper.<commit_after>package main

import (
	"database/sql"
	"fmt"
	"log"
	"os"

	_ "github.com/denisenkom/go-mssqldb"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	"github.com/spf13/pflag"

	"github.com/go-testfixtures/testfixtures/v3"
)

var version = "master"

func main() {
	log.SetFlags(0)
	log.SetOutput(os.Stderr)

	var (
		versionFlag           bool
		dialect               string
		connString            string
		dir                   string
		files                 []string
		paths                 []string
		useAlterContraint     bool
		skipResetSequences    bool
		resetSequencesTo      int64
		skipTestDatabaseCheck bool
		dumperFlag            bool
		dumperTables          []string
		dumperDir             string
	)

	pflag.BoolVar(&versionFlag, "version", false, "show testfixtures version")
	pflag.StringVarP(&dialect, "dialect", "d", "", "which database system you're using (postgres, timescaledb, mysql, mariadb, sqlite or sqlserver)")
	pflag.StringVarP(&connString, "conn", "c", "", "a database connection string")
	pflag.StringVarP(&dir, "dir", "D", "", "a directory of YAML fixtures to load")
	pflag.StringSliceVarP(&files, "files", "f", nil, "a list of YAML files to load")
	pflag.StringSliceVarP(&paths, "paths", "p", nil, "a list of fixture paths to load (directory or file)")
	pflag.BoolVar(&useAlterContraint, "alter-constraint", false, "use ALTER CONSTRAINT to disable referential integrity (PostgreSQL only)")
	pflag.BoolVar(&skipResetSequences, "no-reset-sequences", false, "skip reset of sequences after loading (PostgreSQL only)")
	pflag.Int64Var(&resetSequencesTo, "reset-sequences-to", 0, "sets the number sequences will be reset after loading fixtures (PostgreSQL only, defaults to 10000)")
	pflag.BoolVar(&skipTestDatabaseCheck, "dangerous-no-test-database-check", false, `skips check for "test" in database name (use with caution)`)
	pflag.BoolVar(&dumperFlag, "dumper", false, "dumping fixtures from the database into a directory")
	pflag.StringSliceVarP(&dumperTables, "dumper-tables", "t", nil, "a list of table to choose which tables you want to dump")
	pflag.StringVarP(&dumperDir, "dumper-dir", "s", "", "sets the directory where the dumping fixtures files will be created")
	pflag.Parse()

	if versionFlag {
		log.Printf("testfixtures version: %s", version)
		return
	}

	if dialect == "" && connString == "" {
		log.Fatal("testfixtures: both --dialect (-d) and --conn (-c) are required")
		return
	}
	if dumperFlag && dumperDir == "" {
		log.Fatalf("testfixtures: if use dumper, --dumper-dir (-s) are required")
	}
	if !dumperFlag && dir == "" && len(files) == 0 && len(paths) == 0 {
		log.Fatal("testfixtures: either --dir (-D) or --files (-f) or --paths (-p) need to be given")
		return
	}

	dialect, err := getDialect(dialect)
	if err != nil {
		log.Fatal(err)
		return
	}

	db, err := sql.Open(dialect, connString)
	if err != nil {
		log.Fatalf("testfixtures: could not connect to database: %v", err)
		return
	}
	defer db.Close()

	if err := db.Ping(); err != nil {
		log.Fatalf("testfixtures: could not ping database: %v", err)
		return
	}

	if dumperFlag {
		dumperOptions := []func(*testfixtures.Dumper) error{
			testfixtures.DumpDatabase(db),
			testfixtures.DumpDialect(dialect),
			testfixtures.DumpDirectory(dumperDir),
		}
		if len(dumperTables) > 0 {
			dumperOptions = append(dumperOptions, testfixtures.DumpTables(dumperTables...))
		}
		dumper, err := testfixtures.NewDumper(dumperOptions...)
		if err != nil {
			log.Fatal(err)
		}
		if err := dumper.Dump(); err != nil {
			log.Fatal(err)
		}
		log.Printf("testfixtures: dumped fixtures file has been created successfully")
		return
	}

	options := []func(*testfixtures.Loader) error{
		testfixtures.Database(db),
		testfixtures.Dialect(dialect),
	}
	if dir != "" {
		options = append(options, testfixtures.Directory(dir))
	}
	if len(files) > 0 {
		options = append(options, testfixtures.Files(files...))
	}
	if len(paths) > 0 {
		options = append(options, testfixtures.Paths(paths...))
	}
	if useAlterContraint {
		options = append(options, testfixtures.UseAlterConstraint())
	}
	if skipResetSequences {
		options = append(options, testfixtures.SkipResetSequences())
	}
	if resetSequencesTo > 0 {
		options = append(options, testfixtures.ResetSequencesTo(resetSequencesTo))
	}
	if skipTestDatabaseCheck {
		options = append(options, testfixtures.DangerousSkipTestDatabaseCheck())
	}

	loader, err := testfixtures.New(options...)
	if err != nil {
		log.Fatal(err)
	}
	if err := loader.Load(); err != nil {
		log.Fatal(err)
	}
	log.Printf("testfixtures: fixtures loaded successfully")
}

func getDialect(dialect string) (string, error) {
	switch dialect {
	case "postgres", "postgresql", "timescaledb":
		return "postgres", nil
	case "mysql", "mariadb":
		return "mysql", nil
	case "sqlite", "sqlite3":
		if !isSQLiteSupported() {
			return "", fmt.Errorf("testfixtures: SQLite is not supported in this build")
		}
		return "sqlite3", nil
	case "mssql", "sqlserver":
		return "sqlserver", nil
	default:
		return "", fmt.Errorf(`testfixtures: unrecognized dialect "%s"`, dialect)
	}
}

func isSQLiteSupported() bool {
	for _, d := range sql.Drivers() {
		if d == "sqlite3" {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "net/http"
    "os"
    "reflect"
    "strings"

    "github.com/op/go-logging"
    "github.com/gorilla/websocket"
    "github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

// Get Consumer key/secret and Access token/secret from enviroment variables
var consumerKey = os.Getenv("TWITTER_CONSUMER_KEY")
var consumerSecret = os.Getenv("TWITTER_CONSUMER_SECRET")
var accessToken = os.Getenv("TWITTER_ACCESS_TOKEN")
var accessSecret = os.Getenv("TWITTER_ACCESS_SECRET")

var config = oauth1.NewConfig(consumerKey, consumerSecret)
var token = oauth1.NewToken(accessToken, accessSecret)

// Auto Oauth1
var httpClient = config.Client(oauth1.NoContext, token)

// Twitter Client
var client = twitter.NewClient(httpClient)

// Websocket Upgrade buffer
// Allow any domain can access stream API
var upgrader = websocket.Upgrader {
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool { return true },
}

// Store the tweet filter words
var filters = []string{};

// Store the WebSocket connections
var conns = make(map[string]*websocket.Conn)

var log = logging.MustGetLogger("Twitter Tweeds Log")

// Twitter stream
var globalTwitterStream *twitter.Stream

// Start the app then Listen and Serve
func main() {
//     port := os.Getenv("PORT")

//     if port == "" {
// 		log.Fatal("$PORT must be set")
// 	}

    http.HandleFunc("/", home)
    http.HandleFunc("/twitter/stream", twitterStream)
    fmt.Println("TweetDeck started and ready to rock");
    http.ListenAndServe(":8080",nil)
}

// Serve home page
func home(res http.ResponseWriter, req *http.Request) {
    res.Header().Set("Access-Control-Allow-Origin", "*")
    http.ServeFile(res,req, "./statis/index.html")
}

// Serve Websocket Twitter stream
func twitterStream(res http.ResponseWriter, req *http.Request) {
    log.Info("WS Twitter Request: ",req.URL.Query())
    filter := req.URL.Query()["filter"]
    if filter != nil {

        // Upgrades the http server connection to the websocket protocol
        conn, _ := upgrader.Upgrade(res, req, nil)

        filters = append(filters,filter[0]);
        conns[filter[0]] = conn;
        printConnsFiltes()

        // Stop the twitter stream in order to stream again with new filter
        stopGlobTwitterStream()

    	// Read only tweet from twitter stream
    	demux := twitter.NewSwitchDemux()
    	demux.Tweet = func(tweet *twitter.Tweet) {
    	    text := tweet.Text

    	    for _, v := range filters {
    	       if strings.Contains(text,v) && conns[v] != nil {
    	           wsWriter(conns[v],text)
    	       }
    	    }
    	}

    	fmt.Println("Starting Stream...")

    	// Filter
    	filterParams := &twitter.StreamFilterParams{
    		Track:         filters,
    		StallWarnings: twitter.Bool(true),
    	}

    	stream, err := client.Streams.Filter(filterParams)

    	globalTwitterStream = stream

    	fmt.Println(reflect.TypeOf(stream))
    	if err != nil {
    		log.Error(err)
    	}

        go wsReader(conn,filter[0],stream)      // Read messages from the websocket or stop websoket and twitter stream
    	go demux.HandleChan(stream.Messages)    // Receive twitter messages until stream quits

    } else {
        fmt.Fprintf(res, "Warning: Filter param is required to stream the data from Twitter")
    }
}

func stopGlobTwitterStream() {
    if globalTwitterStream != nil {
            globalTwitterStream.Stop()
    }
}

func printConnsFiltes() {
    fmt.Println("Filters:",filters)
    fmt.Println("Connections:",conns)
}

// Read messages from the websocket
func wsReader(conn *websocket.Conn,filter string,stream *twitter.Stream) {
    for {
        _, _, err := conn.ReadMessage()

        // Close the websocket connection
        if err != nil {
            // Close connection
            conn.Close()
             // Remove connections from connection list
            delete(conns, filter)

            for index, value := range filters {
                if value == filter {
                    // Remove filtered word from filter slice
                    filters = append(filters[:index], filters[index+1:]...)
                    break
                }
            }

            log.Info("wsReader/WS Connection Closed: "+filter)
            break
        } else {
            log.Error(err)
        }
    }
}

// Write messages to the websocket
func wsWriter(conn *websocket.Conn,filter string) {
    conn.WriteJSON(tweetStruct{
        Tweet: filter,
    })
}

type tweetStruct struct {
    Tweet string
}<commit_msg>Fix index.html path<commit_after>package main

import (
    "fmt"
    "net/http"
    "os"
    "reflect"
    "strings"

    "github.com/op/go-logging"
    "github.com/gorilla/websocket"
    "github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

// Get Consumer key/secret and Access token/secret from enviroment variables
var consumerKey = os.Getenv("TWITTER_CONSUMER_KEY")
var consumerSecret = os.Getenv("TWITTER_CONSUMER_SECRET")
var accessToken = os.Getenv("TWITTER_ACCESS_TOKEN")
var accessSecret = os.Getenv("TWITTER_ACCESS_SECRET")

var config = oauth1.NewConfig(consumerKey, consumerSecret)
var token = oauth1.NewToken(accessToken, accessSecret)

// Auto Oauth1
var httpClient = config.Client(oauth1.NoContext, token)

// Twitter Client
var client = twitter.NewClient(httpClient)

// Websocket Upgrade buffer
// Allow any domain can access stream API
var upgrader = websocket.Upgrader {
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool { return true },
}

// Store the tweet filter words
var filters = []string{};

// Store the WebSocket connections
var conns = make(map[string]*websocket.Conn)

var log = logging.MustGetLogger("Twitter Tweeds Log")

// Twitter stream
var globalTwitterStream *twitter.Stream

// Start the app then Listen and Serve
func main() {
    port := os.Getenv("PORT")

    if port == "" {
		log.Fatal("$PORT must be set")
	}

    http.HandleFunc("/", home)
    http.HandleFunc("/twitter/stream", twitterStream)
    fmt.Println("TweetDeck started and ready to rock");
    http.ListenAndServe(port,nil)
}

// Serve home page
func home(res http.ResponseWriter, req *http.Request) {
    res.Header().Set("Access-Control-Allow-Origin", "*")
    http.ServeFile(res,req, "./static/index.html")
}

// Serve Websocket Twitter stream
func twitterStream(res http.ResponseWriter, req *http.Request) {
    log.Info("WS Twitter Request: ",req.URL.Query())
    filter := req.URL.Query()["filter"]
    if filter != nil {

        // Upgrades the http server connection to the websocket protocol
        conn, _ := upgrader.Upgrade(res, req, nil)

        filters = append(filters,filter[0]);
        conns[filter[0]] = conn;
        printConnsFiltes()

        // Stop the twitter stream in order to stream again with new filter
        stopGlobTwitterStream()

    	// Read only tweet from twitter stream
    	demux := twitter.NewSwitchDemux()
    	demux.Tweet = func(tweet *twitter.Tweet) {
    	    text := tweet.Text

    	    for _, v := range filters {
    	       if strings.Contains(text,v) && conns[v] != nil {
    	           wsWriter(conns[v],text)
    	       }
    	    }
    	}

    	fmt.Println("Starting Stream...")

    	// Filter
    	filterParams := &twitter.StreamFilterParams{
    		Track:         filters,
    		StallWarnings: twitter.Bool(true),
    	}

    	stream, err := client.Streams.Filter(filterParams)

    	globalTwitterStream = stream

    	fmt.Println(reflect.TypeOf(stream))
    	if err != nil {
    		log.Error(err)
    	}

        go wsReader(conn,filter[0],stream)      // Read messages from the websocket or stop websoket and twitter stream
    	go demux.HandleChan(stream.Messages)    // Receive twitter messages until stream quits

    } else {
        fmt.Fprintf(res, "Warning: Filter param is required to stream the data from Twitter")
    }
}

func stopGlobTwitterStream() {
    if globalTwitterStream != nil {
            globalTwitterStream.Stop()
    }
}

func printConnsFiltes() {
    fmt.Println("Filters:",filters)
    fmt.Println("Connections:",conns)
}

// Read messages from the websocket
func wsReader(conn *websocket.Conn,filter string,stream *twitter.Stream) {
    for {
        _, _, err := conn.ReadMessage()

        // Close the websocket connection
        if err != nil {
            // Close connection
            conn.Close()
             // Remove connections from connection list
            delete(conns, filter)

            for index, value := range filters {
                if value == filter {
                    // Remove filtered word from filter slice
                    filters = append(filters[:index], filters[index+1:]...)
                    break
                }
            }

            log.Info("wsReader/WS Connection Closed: "+filter)
            break
        } else {
            log.Error(err)
        }
    }
}

// Write messages to the websocket
func wsWriter(conn *websocket.Conn,filter string) {
    conn.WriteJSON(tweetStruct{
        Tweet: filter,
    })
}

type tweetStruct struct {
    Tweet string
}<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package main

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/golang/glog"
	flag "github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/api/errors"

	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes/scheme"
	k8coresv1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"
	"k8s.io/client-go/util/cert"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/certificates"
	"kubevirt.io/kubevirt/pkg/certificates/triple"
	"kubevirt.io/kubevirt/pkg/controller"
	inotifyinformer "kubevirt.io/kubevirt/pkg/inotify-informer"
	_ "kubevirt.io/kubevirt/pkg/monitoring/client/prometheus"    // import for prometheus metrics
	_ "kubevirt.io/kubevirt/pkg/monitoring/reflector/prometheus" // import for prometheus metrics
	promvm "kubevirt.io/kubevirt/pkg/monitoring/vms/prometheus"  // import for prometheus metrics
	_ "kubevirt.io/kubevirt/pkg/monitoring/workqueue/prometheus" // import for prometheus metrics
	"kubevirt.io/kubevirt/pkg/service"
	"kubevirt.io/kubevirt/pkg/util"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	virthandler "kubevirt.io/kubevirt/pkg/virt-handler"
	virtcache "kubevirt.io/kubevirt/pkg/virt-handler/cache"
	virtlauncher "kubevirt.io/kubevirt/pkg/virt-launcher"
	virt_api "kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

const (
	defaultWatchdogTimeout = 15 * time.Second

	// Default port that virt-handler listens on.
	defaultPort = 8185

	// Default address that virt-handler listens on.
	defaultHost = "0.0.0.0"

	hostOverride = ""

	podIpAddress = ""

	virtShareDir = "/var/run/kubevirt"

	// This value is derived from default MaxPods in Kubelet Config
	maxDevices = 110

	clientCertBytesValue  = "client-cert-bytes"
	clientKeyBytesValue   = "client-key-bytes"
	signingCertBytesValue = "signing-cert-bytes"

	// selfsigned cert secret name
	virtHandlerCertSecretName = "kubevirt-virt-handler-certs"
	maxRequestsInFlight       = 3
)

type virtHandlerApp struct {
	service.ServiceListen
	HostOverride            string
	PodIpAddress            string
	VirtShareDir            string
	WatchdogTimeoutDuration time.Duration
	MaxDevices              int
	MaxRequestsInFlight     int

	signingCertBytes []byte
	clientCertBytes  []byte
	clientKeyBytes   []byte

	virtCli   kubecli.KubevirtClient
	namespace string

	migrationTLSConfig *tls.Config
}

var _ service.Service = &virtHandlerApp{}

func (app *virtHandlerApp) getSelfSignedCert() error {
	var ok bool

	caKeyPair, _ := triple.NewCA("kubevirt.io")
	clientKeyPair, _ := triple.NewClientKeyPair(caKeyPair,
		"kubevirt.io:system:node:virt-handler",
		nil,
	)

	secret := &k8sv1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      virtHandlerCertSecretName,
			Namespace: app.namespace,
			Labels: map[string]string{
				v1.AppLabel: "virt-api-aggregator",
			},
		},
		Type: "Opaque",
		Data: map[string][]byte{
			clientCertBytesValue:  cert.EncodeCertPEM(clientKeyPair.Cert),
			clientKeyBytesValue:   cert.EncodePrivateKeyPEM(clientKeyPair.Key),
			signingCertBytesValue: cert.EncodeCertPEM(caKeyPair.Cert),
		},
	}
	_, err := app.virtCli.CoreV1().Secrets(app.namespace).Create(secret)
	if errors.IsAlreadyExists(err) {
		secret, err = app.virtCli.CoreV1().Secrets(app.namespace).Get(virtHandlerCertSecretName, metav1.GetOptions{})
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	// retrieve self signed cert info from secret
	app.clientCertBytes, ok = secret.Data[clientCertBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", clientCertBytesValue, virtHandlerCertSecretName)
	}
	app.clientKeyBytes, ok = secret.Data[clientKeyBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", clientKeyBytesValue, virtHandlerCertSecretName)
	}
	app.signingCertBytes, ok = secret.Data[signingCertBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", signingCertBytesValue, virtHandlerCertSecretName)
	}
	return nil
}

func (app *virtHandlerApp) Run() {
	// HostOverride should default to os.Hostname(), to make sure we handle errors ensure it here.
	if app.HostOverride == "" {
		defaultHostName, err := os.Hostname()
		if err != nil {
			panic(err)
		}
		app.HostOverride = defaultHostName
	}

	if app.PodIpAddress == "" {
		panic(fmt.Errorf("no pod ip detected"))
	}

	logger := log.Log
	logger.V(1).Level(log.INFO).Log("hostname", app.HostOverride)

	// Create event recorder
	var err error
	app.virtCli, err = kubecli.GetKubevirtClient()
	if err != nil {
		panic(err)
	}
	broadcaster := record.NewBroadcaster()
	broadcaster.StartRecordingToSink(&k8coresv1.EventSinkImpl{Interface: app.virtCli.CoreV1().Events(k8sv1.NamespaceAll)})
	// Scheme is used to create an ObjectReference from an Object (e.g. VirtualMachineInstance) during Event creation
	recorder := broadcaster.NewRecorder(scheme.Scheme, k8sv1.EventSource{Component: "virt-handler", Host: app.HostOverride})

	vmiSourceLabel, err := labels.Parse(fmt.Sprintf(v1.NodeNameLabel+" in (%s)", app.HostOverride))
	if err != nil {
		panic(err)
	}
	vmiTargetLabel, err := labels.Parse(fmt.Sprintf(v1.MigrationTargetNodeNameLabel+" in (%s)", app.HostOverride))
	if err != nil {
		panic(err)
	}

	// Wire VirtualMachineInstance controller

	vmSourceSharedInformer := cache.NewSharedIndexInformer(
		controller.NewListWatchFromClient(app.virtCli.RestClient(), "virtualmachineinstances", k8sv1.NamespaceAll, fields.Everything(), vmiSourceLabel),
		&v1.VirtualMachineInstance{},
		0,
		cache.Indexers{},
	)

	vmTargetSharedInformer := cache.NewSharedIndexInformer(
		controller.NewListWatchFromClient(app.virtCli.RestClient(), "virtualmachineinstances", k8sv1.NamespaceAll, fields.Everything(), vmiTargetLabel),
		&v1.VirtualMachineInstance{},
		0,
		cache.Indexers{},
	)

	// Wire Domain controller
	domainSharedInformer, err := virtcache.NewSharedInformer(app.VirtShareDir, int(app.WatchdogTimeoutDuration.Seconds()), recorder, vmSourceSharedInformer.GetStore())
	if err != nil {
		panic(err)
	}

	virtlauncher.InitializeSharedDirectories(app.VirtShareDir)

	app.namespace, err = util.GetNamespace()
	if err != nil {
		glog.Fatalf("Error searching for namespace: %v", err)
	}

	if err := app.getSelfSignedCert(); err != nil {
		glog.Fatalf("Error loading self signed certificates: %v", err)
	}

	if err := app.setupTLS(); err != nil {
		glog.Fatalf("Error constructing migration tls config: %v", err)
	}

	factory := controller.NewKubeInformerFactory(app.virtCli.RestClient(), app.virtCli, app.namespace)

	gracefulShutdownInformer := cache.NewSharedIndexInformer(
		inotifyinformer.NewFileListWatchFromClient(
			virtlauncher.GracefulShutdownTriggerDir(app.VirtShareDir)),
		&virt_api.Domain{},
		0,
		cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc})

	vmController := virthandler.NewController(
		recorder,
		app.virtCli,
		app.HostOverride,
		app.PodIpAddress,
		app.VirtShareDir,
		vmSourceSharedInformer,
		vmTargetSharedInformer,
		domainSharedInformer,
		gracefulShutdownInformer,
		int(app.WatchdogTimeoutDuration.Seconds()),
		app.MaxDevices,
		virtconfig.NewClusterConfig(factory.ConfigMap(), app.namespace),
		app.migrationTLSConfig,
	)

	certsDirectory, err := ioutil.TempDir("", "certsdir")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(certsDirectory)

	certStore, err := certificates.GenerateSelfSignedCert(certsDirectory, "virt-handler", app.namespace)
	if err != nil {
		glog.Fatalf("unable to generate certificates: %v", err)
	}

	promvm.SetupCollector(app.virtCli, app.VirtShareDir, app.HostOverride)

	// Bootstrapping. From here on the startup order matters
	stop := make(chan struct{})
	defer close(stop)
	factory.Start(stop)
	cache.WaitForCacheSync(stop, factory.ConfigMap().HasSynced)

	go vmController.Run(3, stop)

	logger.V(1).Infof("metrics: max concurrent requests=%d", app.MaxRequestsInFlight)
	http.Handle("/metrics", promvm.Handler(app.MaxRequestsInFlight))

	err = http.ListenAndServeTLS(app.ServiceListen.Address(), certStore.CurrentPath(), certStore.CurrentPath(), nil)
	if err != nil {
		log.Log.Reason(err).Error("Serving prometheus failed.")
		panic(err)
	}
}

func (app *virtHandlerApp) AddFlags() {
	app.InitFlags()

	app.BindAddress = defaultHost
	app.Port = defaultPort

	app.AddCommonFlags()

	flag.StringVar(&app.HostOverride, "hostname-override", hostOverride,
		"Name under which the node is registered in Kubernetes, where this virt-handler instance is running on")

	flag.StringVar(&app.PodIpAddress, "pod-ip-address", podIpAddress,
		"The pod ip address")

	flag.StringVar(&app.VirtShareDir, "kubevirt-share-dir", virtShareDir,
		"Shared directory between virt-handler and virt-launcher")

	flag.DurationVar(&app.WatchdogTimeoutDuration, "watchdog-timeout", defaultWatchdogTimeout,
		"Watchdog file timeout")

	// TODO: the Device Plugin API does not allow for infinitely available (shared) devices
	// so the current approach is to register an arbitrary number.
	// This should be deprecated if the API allows for shared resources in the future
	flag.IntVar(&app.MaxDevices, "max-devices", maxDevices,
		"Number of devices to register with Kubernetes device plugin framework")

	flag.IntVar(&app.MaxRequestsInFlight, "max-metric-requests", maxRequestsInFlight,
		"Number of concurrent requests to the metrics endpoint")
}

func (app *virtHandlerApp) setupTLS() error {

	clientCert, err := tls.X509KeyPair(app.clientCertBytes, app.clientKeyBytes)
	if err != nil {
		return err
	}

	caCert, err := cert.ParseCertsPEM(app.signingCertBytes)
	if err != nil {
		return err
	}

	certPool := x509.NewCertPool()

	for _, crt := range caCert {
		certPool.AddCert(crt)
	}

	app.migrationTLSConfig = &tls.Config{
		MinVersion: tls.VersionTLS12,
		ClientCAs:  certPool,
		GetClientCertificate: func(info *tls.CertificateRequestInfo) (certificate *tls.Certificate, e error) {
			return &clientCert, nil
		},
		GetCertificate: func(info *tls.ClientHelloInfo) (i *tls.Certificate, e error) {
			return &clientCert, nil
		},
		// Neither the client nor the server should validate anything itself, `VerifyPeerCertificate` is still executed
		InsecureSkipVerify: true,
		// XXX: We need to verify the cert ourselves because we don't have DNS or IP on the certs at the moment
		VerifyPeerCertificate: func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {

			// impossible with RequireAnyClientCert
			if len(rawCerts) == 0 {
				return fmt.Errorf("no client certificate provided.")
			}

			c, err := x509.ParseCertificate(rawCerts[0])
			if err != nil {
				return fmt.Errorf("failed to parse peer certificate: %v", err)
			}
			_, err = c.Verify(x509.VerifyOptions{
				Roots:     certPool,
				KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
			})

			if err != nil {
				return fmt.Errorf("could not verify peer certificate: %v", err)
			}
			return nil
		},
		ClientAuth: tls.RequireAndVerifyClientCert,
	}

	return nil
}

func main() {
	app := &virtHandlerApp{}
	service.Setup(app)
	log.InitializeLogging("virt-handler")
	app.Run()
}
<commit_msg>Increase the number of workers<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package main

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/golang/glog"
	flag "github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/api/errors"

	k8sv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes/scheme"
	k8coresv1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"
	"k8s.io/client-go/util/cert"

	v1 "kubevirt.io/client-go/api/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/certificates"
	"kubevirt.io/kubevirt/pkg/certificates/triple"
	"kubevirt.io/kubevirt/pkg/controller"
	inotifyinformer "kubevirt.io/kubevirt/pkg/inotify-informer"
	_ "kubevirt.io/kubevirt/pkg/monitoring/client/prometheus"    // import for prometheus metrics
	_ "kubevirt.io/kubevirt/pkg/monitoring/reflector/prometheus" // import for prometheus metrics
	promvm "kubevirt.io/kubevirt/pkg/monitoring/vms/prometheus"  // import for prometheus metrics
	_ "kubevirt.io/kubevirt/pkg/monitoring/workqueue/prometheus" // import for prometheus metrics
	"kubevirt.io/kubevirt/pkg/service"
	"kubevirt.io/kubevirt/pkg/util"
	virtconfig "kubevirt.io/kubevirt/pkg/virt-config"
	virthandler "kubevirt.io/kubevirt/pkg/virt-handler"
	virtcache "kubevirt.io/kubevirt/pkg/virt-handler/cache"
	virtlauncher "kubevirt.io/kubevirt/pkg/virt-launcher"
	virt_api "kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

const (
	defaultWatchdogTimeout = 15 * time.Second

	// Default port that virt-handler listens on.
	defaultPort = 8185

	// Default address that virt-handler listens on.
	defaultHost = "0.0.0.0"

	hostOverride = ""

	podIpAddress = ""

	virtShareDir = "/var/run/kubevirt"

	// This value is derived from default MaxPods in Kubelet Config
	maxDevices = 110

	clientCertBytesValue  = "client-cert-bytes"
	clientKeyBytesValue   = "client-key-bytes"
	signingCertBytesValue = "signing-cert-bytes"

	// selfsigned cert secret name
	virtHandlerCertSecretName = "kubevirt-virt-handler-certs"
	maxRequestsInFlight       = 3
)

type virtHandlerApp struct {
	service.ServiceListen
	HostOverride            string
	PodIpAddress            string
	VirtShareDir            string
	WatchdogTimeoutDuration time.Duration
	MaxDevices              int
	MaxRequestsInFlight     int

	signingCertBytes []byte
	clientCertBytes  []byte
	clientKeyBytes   []byte

	virtCli   kubecli.KubevirtClient
	namespace string

	migrationTLSConfig *tls.Config
}

var _ service.Service = &virtHandlerApp{}

func (app *virtHandlerApp) getSelfSignedCert() error {
	var ok bool

	caKeyPair, _ := triple.NewCA("kubevirt.io")
	clientKeyPair, _ := triple.NewClientKeyPair(caKeyPair,
		"kubevirt.io:system:node:virt-handler",
		nil,
	)

	secret := &k8sv1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      virtHandlerCertSecretName,
			Namespace: app.namespace,
			Labels: map[string]string{
				v1.AppLabel: "virt-api-aggregator",
			},
		},
		Type: "Opaque",
		Data: map[string][]byte{
			clientCertBytesValue:  cert.EncodeCertPEM(clientKeyPair.Cert),
			clientKeyBytesValue:   cert.EncodePrivateKeyPEM(clientKeyPair.Key),
			signingCertBytesValue: cert.EncodeCertPEM(caKeyPair.Cert),
		},
	}
	_, err := app.virtCli.CoreV1().Secrets(app.namespace).Create(secret)
	if errors.IsAlreadyExists(err) {
		secret, err = app.virtCli.CoreV1().Secrets(app.namespace).Get(virtHandlerCertSecretName, metav1.GetOptions{})
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	// retrieve self signed cert info from secret
	app.clientCertBytes, ok = secret.Data[clientCertBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", clientCertBytesValue, virtHandlerCertSecretName)
	}
	app.clientKeyBytes, ok = secret.Data[clientKeyBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", clientKeyBytesValue, virtHandlerCertSecretName)
	}
	app.signingCertBytes, ok = secret.Data[signingCertBytesValue]
	if !ok {
		return fmt.Errorf("%s value not found in %s virt-api secret", signingCertBytesValue, virtHandlerCertSecretName)
	}
	return nil
}

func (app *virtHandlerApp) Run() {
	// HostOverride should default to os.Hostname(), to make sure we handle errors ensure it here.
	if app.HostOverride == "" {
		defaultHostName, err := os.Hostname()
		if err != nil {
			panic(err)
		}
		app.HostOverride = defaultHostName
	}

	if app.PodIpAddress == "" {
		panic(fmt.Errorf("no pod ip detected"))
	}

	logger := log.Log
	logger.V(1).Level(log.INFO).Log("hostname", app.HostOverride)

	// Create event recorder
	var err error
	app.virtCli, err = kubecli.GetKubevirtClient()
	if err != nil {
		panic(err)
	}
	broadcaster := record.NewBroadcaster()
	broadcaster.StartRecordingToSink(&k8coresv1.EventSinkImpl{Interface: app.virtCli.CoreV1().Events(k8sv1.NamespaceAll)})
	// Scheme is used to create an ObjectReference from an Object (e.g. VirtualMachineInstance) during Event creation
	recorder := broadcaster.NewRecorder(scheme.Scheme, k8sv1.EventSource{Component: "virt-handler", Host: app.HostOverride})

	vmiSourceLabel, err := labels.Parse(fmt.Sprintf(v1.NodeNameLabel+" in (%s)", app.HostOverride))
	if err != nil {
		panic(err)
	}
	vmiTargetLabel, err := labels.Parse(fmt.Sprintf(v1.MigrationTargetNodeNameLabel+" in (%s)", app.HostOverride))
	if err != nil {
		panic(err)
	}

	// Wire VirtualMachineInstance controller

	vmSourceSharedInformer := cache.NewSharedIndexInformer(
		controller.NewListWatchFromClient(app.virtCli.RestClient(), "virtualmachineinstances", k8sv1.NamespaceAll, fields.Everything(), vmiSourceLabel),
		&v1.VirtualMachineInstance{},
		0,
		cache.Indexers{},
	)

	vmTargetSharedInformer := cache.NewSharedIndexInformer(
		controller.NewListWatchFromClient(app.virtCli.RestClient(), "virtualmachineinstances", k8sv1.NamespaceAll, fields.Everything(), vmiTargetLabel),
		&v1.VirtualMachineInstance{},
		0,
		cache.Indexers{},
	)

	// Wire Domain controller
	domainSharedInformer, err := virtcache.NewSharedInformer(app.VirtShareDir, int(app.WatchdogTimeoutDuration.Seconds()), recorder, vmSourceSharedInformer.GetStore())
	if err != nil {
		panic(err)
	}

	virtlauncher.InitializeSharedDirectories(app.VirtShareDir)

	app.namespace, err = util.GetNamespace()
	if err != nil {
		glog.Fatalf("Error searching for namespace: %v", err)
	}

	if err := app.getSelfSignedCert(); err != nil {
		glog.Fatalf("Error loading self signed certificates: %v", err)
	}

	if err := app.setupTLS(); err != nil {
		glog.Fatalf("Error constructing migration tls config: %v", err)
	}

	factory := controller.NewKubeInformerFactory(app.virtCli.RestClient(), app.virtCli, app.namespace)

	gracefulShutdownInformer := cache.NewSharedIndexInformer(
		inotifyinformer.NewFileListWatchFromClient(
			virtlauncher.GracefulShutdownTriggerDir(app.VirtShareDir)),
		&virt_api.Domain{},
		0,
		cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc})

	vmController := virthandler.NewController(
		recorder,
		app.virtCli,
		app.HostOverride,
		app.PodIpAddress,
		app.VirtShareDir,
		vmSourceSharedInformer,
		vmTargetSharedInformer,
		domainSharedInformer,
		gracefulShutdownInformer,
		int(app.WatchdogTimeoutDuration.Seconds()),
		app.MaxDevices,
		virtconfig.NewClusterConfig(factory.ConfigMap(), app.namespace),
		app.migrationTLSConfig,
	)

	certsDirectory, err := ioutil.TempDir("", "certsdir")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(certsDirectory)

	certStore, err := certificates.GenerateSelfSignedCert(certsDirectory, "virt-handler", app.namespace)
	if err != nil {
		glog.Fatalf("unable to generate certificates: %v", err)
	}

	promvm.SetupCollector(app.virtCli, app.VirtShareDir, app.HostOverride)

	// Bootstrapping. From here on the startup order matters
	stop := make(chan struct{})
	defer close(stop)
	factory.Start(stop)
	cache.WaitForCacheSync(stop, factory.ConfigMap().HasSynced)

	go vmController.Run(10, stop)

	logger.V(1).Infof("metrics: max concurrent requests=%d", app.MaxRequestsInFlight)
	http.Handle("/metrics", promvm.Handler(app.MaxRequestsInFlight))

	err = http.ListenAndServeTLS(app.ServiceListen.Address(), certStore.CurrentPath(), certStore.CurrentPath(), nil)
	if err != nil {
		log.Log.Reason(err).Error("Serving prometheus failed.")
		panic(err)
	}
}

func (app *virtHandlerApp) AddFlags() {
	app.InitFlags()

	app.BindAddress = defaultHost
	app.Port = defaultPort

	app.AddCommonFlags()

	flag.StringVar(&app.HostOverride, "hostname-override", hostOverride,
		"Name under which the node is registered in Kubernetes, where this virt-handler instance is running on")

	flag.StringVar(&app.PodIpAddress, "pod-ip-address", podIpAddress,
		"The pod ip address")

	flag.StringVar(&app.VirtShareDir, "kubevirt-share-dir", virtShareDir,
		"Shared directory between virt-handler and virt-launcher")

	flag.DurationVar(&app.WatchdogTimeoutDuration, "watchdog-timeout", defaultWatchdogTimeout,
		"Watchdog file timeout")

	// TODO: the Device Plugin API does not allow for infinitely available (shared) devices
	// so the current approach is to register an arbitrary number.
	// This should be deprecated if the API allows for shared resources in the future
	flag.IntVar(&app.MaxDevices, "max-devices", maxDevices,
		"Number of devices to register with Kubernetes device plugin framework")

	flag.IntVar(&app.MaxRequestsInFlight, "max-metric-requests", maxRequestsInFlight,
		"Number of concurrent requests to the metrics endpoint")
}

func (app *virtHandlerApp) setupTLS() error {

	clientCert, err := tls.X509KeyPair(app.clientCertBytes, app.clientKeyBytes)
	if err != nil {
		return err
	}

	caCert, err := cert.ParseCertsPEM(app.signingCertBytes)
	if err != nil {
		return err
	}

	certPool := x509.NewCertPool()

	for _, crt := range caCert {
		certPool.AddCert(crt)
	}

	app.migrationTLSConfig = &tls.Config{
		MinVersion: tls.VersionTLS12,
		ClientCAs:  certPool,
		GetClientCertificate: func(info *tls.CertificateRequestInfo) (certificate *tls.Certificate, e error) {
			return &clientCert, nil
		},
		GetCertificate: func(info *tls.ClientHelloInfo) (i *tls.Certificate, e error) {
			return &clientCert, nil
		},
		// Neither the client nor the server should validate anything itself, `VerifyPeerCertificate` is still executed
		InsecureSkipVerify: true,
		// XXX: We need to verify the cert ourselves because we don't have DNS or IP on the certs at the moment
		VerifyPeerCertificate: func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {

			// impossible with RequireAnyClientCert
			if len(rawCerts) == 0 {
				return fmt.Errorf("no client certificate provided.")
			}

			c, err := x509.ParseCertificate(rawCerts[0])
			if err != nil {
				return fmt.Errorf("failed to parse peer certificate: %v", err)
			}
			_, err = c.Verify(x509.VerifyOptions{
				Roots:     certPool,
				KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
			})

			if err != nil {
				return fmt.Errorf("could not verify peer certificate: %v", err)
			}
			return nil
		},
		ClientAuth: tls.RequireAndVerifyClientCert,
	}

	return nil
}

func main() {
	app := &virtHandlerApp{}
	service.Setup(app)
	log.InitializeLogging("virt-handler")
	app.Run()
}
<|endoftext|>


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

package tcp_test

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"testing"
	"time"

	"github.com/mikioh/tcp"
)

var (
	host = "www.google.com"
	url  = "https://www.google.com/robots.txt"
	tt   *testing.T
)

func TestInfoWithGoogle(t *testing.T) {
	switch {
	case runtime.GOOS == "freebsd":
	case runtime.GOOS == "linux" && runtime.GOARCH == "amd64" || runtime.GOOS == "linux" && runtime.GOARCH == "arm":
	default:
		t.Skipf("not supported on %s/%s", runtime.GOOS, runtime.GOARCH)
	}

	tt = t
	tr := &http.Transport{
		Dial:            dialWithTCPConnMonitor,
		TLSClientConfig: &tls.Config{ServerName: host},
	}
	client := http.Client{Transport: tr}
	resp, err := client.Get(url)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		t.Fatal(err)
	}
}

func dialWithTCPConnMonitor(network, address string) (net.Conn, error) {
	d := net.Dialer{DualStack: true}
	c, err := d.Dial(network, address)
	if err != nil {
		return nil, err
	}
	tc, err := tcp.NewConn(c)
	if err != nil {
		c.Close()
		return nil, err
	}
	go tcpConnMonitor(tc)
	return tc, nil
}

func tcpConnMonitor(c *tcp.Conn) {
	tt.Logf("%v -> %v", c.LocalAddr(), c.RemoteAddr())
	for {
		ti, err := c.Info()
		if err != nil {
			break
		}
		text, err := json.Marshal(ti)
		if err != nil {
			tt.Error(err)
			return
		}
		tt.Log(string(text))
		time.Sleep(20 * time.Millisecond)
	}
}
<commit_msg>tcp: simplify test<commit_after>// Copyright 2014 Mikio Hara. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tcp_test

import (
	"crypto/tls"
	"encoding/json"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"testing"
	"time"

	"github.com/mikioh/tcp"
)

var (
	host = "www.google.com"
	url  = "https://www.google.com/robots.txt"
	tt   *testing.T
)

func TestInfoWithGoogle(t *testing.T) {
	switch {
	case runtime.GOOS == "freebsd":
	case runtime.GOOS == "linux" && (runtime.GOARCH == "amd64" || runtime.GOARCH == "arm"):
	default:
		t.Skipf("not supported on %s/%s", runtime.GOOS, runtime.GOARCH)
	}

	tt = t
	tr := &http.Transport{
		Dial:            dialWithTCPConnMonitor,
		TLSClientConfig: &tls.Config{ServerName: host},
	}
	client := http.Client{Transport: tr}
	resp, err := client.Get(url)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		t.Fatal(err)
	}
}

func dialWithTCPConnMonitor(network, address string) (net.Conn, error) {
	d := net.Dialer{DualStack: true}
	c, err := d.Dial(network, address)
	if err != nil {
		return nil, err
	}
	tc, err := tcp.NewConn(c)
	if err != nil {
		c.Close()
		return nil, err
	}
	go tcpConnMonitor(tc)
	return tc, nil
}

func tcpConnMonitor(c *tcp.Conn) {
	tt.Logf("%v -> %v", c.LocalAddr(), c.RemoteAddr())
	for {
		ti, err := c.Info()
		if err != nil {
			break
		}
		text, err := json.Marshal(ti)
		if err != nil {
			tt.Error(err)
			return
		}
		tt.Log(string(text))
		time.Sleep(20 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 session

import (
	"errors"

	"github.com/admpub/sessions"
	"github.com/webx-top/echo"
)

const CookieMaxAgeKey = `CookieMaxAge`

func Sessions(options *echo.SessionOptions, store sessions.Store) echo.MiddlewareFuncd {
	var newSession func(ctx echo.Context) echo.Sessioner
	if options == nil {
		newSession = func(ctx echo.Context) echo.Sessioner {
			return NewMySession(store, ctx.SessionOptions().Name, ctx)
		}
	} else {
		newSession = func(ctx echo.Context) echo.Sessioner {
			sessionOptions := options.Clone()
			ctx.SetSessionOptions(sessionOptions)
			return NewMySession(store, sessionOptions.Name, ctx)
		}
	}
	return func(h echo.Handler) echo.HandlerFunc {
		return func(c echo.Context) error {
			s := newSession(c)
			c.SetSessioner(s)
			s.SetPreSaveHook(func(c echo.Context) error{
				maxAge, ok := s.Get(CookieMaxAgeKey).(int)
				if !ok {
					return nil
				}
				if maxAge > 3600 && c.CookieOptions().MaxAge != maxAge {
					c.CookieOptions().MaxAge = maxAge
				}
				return nil
			})
			c.AddPreResponseHook(s.Save)
			err := h.Handle(c)
			if e := s.Save(); e != nil {
				if err != nil {
					err = errors.New("Multiple errors:\n1. " + err.Error() + "\n2. " + e.Error())
				} else {
					err = e
				}
				c.Logger().Error(err)
			}
			return err
		}
	}
}

func Middleware(options *echo.SessionOptions) echo.MiddlewareFuncd {
	store := StoreEngine(options)
	return Sessions(options, store)
}
<commit_msg>update<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   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 session

import (
	"errors"

	"github.com/admpub/sessions"
	"github.com/webx-top/echo"
)

const CookieMaxAgeKey = `CookieMaxAge`

func RememberMaxAge(c echo.Context, maxAge int) {
	if maxAge > 0 {
		c.CookieOptions().MaxAge = maxAge
		c.Session().Set(CookieMaxAgeKey, maxAge)
	}
}

func Sessions(options *echo.SessionOptions, store sessions.Store) echo.MiddlewareFuncd {
	var newSession func(ctx echo.Context) echo.Sessioner
	if options == nil {
		newSession = func(ctx echo.Context) echo.Sessioner {
			return NewMySession(store, ctx.SessionOptions().Name, ctx)
		}
	} else {
		newSession = func(ctx echo.Context) echo.Sessioner {
			sessionOptions := options.Clone()
			ctx.SetSessionOptions(sessionOptions)
			return NewMySession(store, sessionOptions.Name, ctx)
		}
	}
	return func(h echo.Handler) echo.HandlerFunc {
		return func(c echo.Context) error {
			s := newSession(c)
			c.SetSessioner(s)
			s.SetPreSaveHook(func(c echo.Context) error{
				maxAge, ok := s.Get(CookieMaxAgeKey).(int)
				if !ok {
					return nil
				}
				if maxAge > 3600 && c.CookieOptions().MaxAge != maxAge {
					c.CookieOptions().MaxAge = maxAge
				}
				return nil
			})
			c.AddPreResponseHook(s.Save)
			err := h.Handle(c)
			if e := s.Save(); e != nil {
				if err != nil {
					err = errors.New("Multiple errors:\n1. " + err.Error() + "\n2. " + e.Error())
				} else {
					err = e
				}
				c.Logger().Error(err)
			}
			return err
		}
	}
}

func Middleware(options *echo.SessionOptions) echo.MiddlewareFuncd {
	store := StoreEngine(options)
	return Sessions(options, store)
}
<|endoftext|>


<commit_before>package cli53

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/route53"
)

type instancesArgs struct {
	name     string
	off      string
	regions  []string
	wait     bool
	ttl      int
	match    string
	internal bool
	aRecord  bool
	dryRun   bool
}

type InstanceRecord struct {
	name  string
	value string
}

func instances(args instancesArgs, config *aws.Config) {
	zone := lookupZone(args.name)
	fmt.Println("Getting DNS records")

	describeInstancesInput := ec2.DescribeInstancesInput{}
	if args.off == "" {
		filter := ec2.Filter{
			Name:   aws.String("instance-state-name"),
			Values: []*string{aws.String("running")},
		}
		describeInstancesInput.Filters = []*ec2.Filter{&filter}
	}

	var reMatch *regexp.Regexp
	if args.match != "" {
		var err error
		reMatch, err = regexp.Compile(args.match)
		if err != nil {
			fatalIfErr(err)
		}
	}

	insts := map[string]*ec2.Instance{}
	for _, region := range args.regions {
		ec2conn := ec2.New(session.New(), config.WithRegion(region))
		output, err := ec2conn.DescribeInstances(&describeInstancesInput)
		fatalIfErr(err)
		for _, r := range output.Reservations {
			for _, i := range r.Instances {
				for _, tag := range i.Tags {
					// limit to instances with a Name tag
					if *tag.Key == "Name" {
						if reMatch != nil && !reMatch.MatchString(*tag.Value) {
							continue
						}
						insts[*tag.Value] = i
						continue
					}
				}
			}
		}
	}

	if len(insts) == 0 {
		fmt.Println("No instances found")
	}

	var rtype string
	if args.aRecord {
		rtype = "A"
	} else {
		rtype = "CNAME"
	}

	suffix := "." + *zone.Name
	suffix = strings.TrimSuffix(suffix, ".")

	upserts := []*route53.Change{}
	for name, instance := range insts {
		var value *string
		if *instance.State.Name != "running" {
			value = &args.off
		} else if args.aRecord {
			if args.internal {
				value = instance.PrivateIpAddress
			} else {
				value = instance.PublicIpAddress
			}
		} else {
			if args.internal {
				value = aws.String(*instance.PrivateDnsName + ".")
			} else {
				value = aws.String(*instance.PublicDnsName + ".")
			}
		}

		// add domain suffix if missing
		dnsname := name
		if !strings.HasSuffix(dnsname, suffix) {
			dnsname += suffix
		}
		rr := route53.ResourceRecord{
			Value: value,
		}
		rrset := route53.ResourceRecordSet{
			Name:            &dnsname,
			TTL:             aws.Int64(int64(args.ttl)),
			Type:            &rtype,
			ResourceRecords: []*route53.ResourceRecord{&rr},
		}
		change := route53.Change{
			Action:            aws.String("UPSERT"),
			ResourceRecordSet: &rrset,
		}
		upserts = append(upserts, &change)
	}

	if args.dryRun {
		fmt.Println("Dry-run, upserts that would be made:")
		for _, upsert := range upserts {
			rr := upsert.ResourceRecordSet
			fmt.Printf("+ %s %s %v\n", *rr.Name, *rr.Type, *rr.ResourceRecords[0].Value)
		}
	} else {
		resp := batchChanges(upserts, []*route53.Change{}, zone)
		fmt.Printf("%d records upserted\n", len(upserts))

		if args.wait && resp != nil {
			waitForChange(resp.ChangeInfo)
		}
	}
}
<commit_msg>Paginate instances listing<commit_after>package cli53

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/route53"
)

type instancesArgs struct {
	name     string
	off      string
	regions  []string
	wait     bool
	ttl      int
	match    string
	internal bool
	aRecord  bool
	dryRun   bool
}

type InstanceRecord struct {
	name  string
	value string
}

func instances(args instancesArgs, config *aws.Config) {
	zone := lookupZone(args.name)
	fmt.Println("Getting DNS records")

	describeInstancesInput := ec2.DescribeInstancesInput{}
	if args.off == "" {
		filter := ec2.Filter{
			Name:   aws.String("instance-state-name"),
			Values: []*string{aws.String("running")},
		}
		describeInstancesInput.Filters = []*ec2.Filter{&filter}
	}

	var reMatch *regexp.Regexp
	if args.match != "" {
		var err error
		reMatch, err = regexp.Compile(args.match)
		if err != nil {
			fatalIfErr(err)
		}
	}

	insts := map[string]*ec2.Instance{}
	for _, region := range args.regions {
		ec2conn := ec2.New(session.New(), config.WithRegion(region))
		for {
			// paginated
			output, err := ec2conn.DescribeInstances(&describeInstancesInput)
			fatalIfErr(err)
			for _, r := range output.Reservations {
				for _, i := range r.Instances {
					for _, tag := range i.Tags {
						// limit to instances with a Name tag
						if *tag.Key == "Name" {
							if reMatch != nil && !reMatch.MatchString(*tag.Value) {
								continue
							}
							insts[*tag.Value] = i
							continue
						}
					}
				}
			}

			if output.NextToken == nil {
				break
			}
			describeInstancesInput.NextToken = output.NextToken
		}
	}

	if len(insts) == 0 {
		fmt.Println("No instances found")
	}

	var rtype string
	if args.aRecord {
		rtype = "A"
	} else {
		rtype = "CNAME"
	}

	suffix := "." + *zone.Name
	suffix = strings.TrimSuffix(suffix, ".")

	upserts := []*route53.Change{}
	for name, instance := range insts {
		var value *string
		if *instance.State.Name != "running" {
			value = &args.off
		} else if args.aRecord {
			if args.internal {
				value = instance.PrivateIpAddress
			} else {
				value = instance.PublicIpAddress
			}
		} else {
			if args.internal {
				value = aws.String(*instance.PrivateDnsName + ".")
			} else {
				value = aws.String(*instance.PublicDnsName + ".")
			}
		}

		// add domain suffix if missing
		dnsname := name
		if !strings.HasSuffix(dnsname, suffix) {
			dnsname += suffix
		}
		rr := route53.ResourceRecord{
			Value: value,
		}
		rrset := route53.ResourceRecordSet{
			Name:            &dnsname,
			TTL:             aws.Int64(int64(args.ttl)),
			Type:            &rtype,
			ResourceRecords: []*route53.ResourceRecord{&rr},
		}
		change := route53.Change{
			Action:            aws.String("UPSERT"),
			ResourceRecordSet: &rrset,
		}
		upserts = append(upserts, &change)
	}

	if args.dryRun {
		fmt.Println("Dry-run, upserts that would be made:")
		for _, upsert := range upserts {
			rr := upsert.ResourceRecordSet
			fmt.Printf("+ %s %s %v\n", *rr.Name, *rr.Type, *rr.ResourceRecords[0].Value)
		}
	} else {
		resp := batchChanges(upserts, []*route53.Change{}, zone)
		fmt.Printf("%d records upserted\n", len(upserts))

		if args.wait && resp != nil {
			waitForChange(resp.ChangeInfo)
		}
	}
}
<|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"time"

	"github.com/bazelbuild/continuous-integration/metrics/clients"
	"github.com/bazelbuild/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

const baseMetricType = "custom.googleapis.com/bazel/ci"

type PlatformLoad struct {
	client  clients.BuildkiteClient
	orgs    []string
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (*PlatformLoad) Type() MetricType {
	return TimeBasedMetric
}

func (*PlatformLoad) RelevantDelta() int {
	return 2 * 24 * 60 * 60 // Two days in seconds
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now(), rows: make([]*loadDataRow, 0)}
	for _, org := range pl.orgs {
		pid := &data.PipelineID{Org: org, Slug: "all"}
		builds, err := pl.client.GetMostRecentBuilds(pid, pl.builds)
		if err != nil {
			return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
		}

		allPlatforms := make(map[string]bool)
		waiting := make(map[string]int)
		running := make(map[string]int)
		for _, build := range builds {
			for _, job := range build.Jobs {
				// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
				platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
				if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
					continue
				}
				allPlatforms[platform] = true
				switch *job.State {
				case "running":
					running[platform] += 1
				case "scheduled", "runnable":
					/*
						State "scheduled" / "runnable" = waiting for a worker to become available
						State "waiting" / "waiting_failed" = waiting for another task to finish

						We're only interested in "scheduled" and "runnable" jobs since they may indicate a shortage of workers.
					*/
					waiting[platform] += 1
				}

			}
		}

		for platform := range allPlatforms {
			row := &loadDataRow{org: org, platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
			result.rows = append(result.rows, row)
		}
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, org VARCHAR(255), platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(org, timestamp, platform));
func CreatePlatformLoad(client clients.BuildkiteClient, builds int, orgs ...string) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"org", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, orgs: orgs, columns: columns, builds: builds}
}

type loadDataRow struct {
	org         string
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []*loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.org, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows)*3)
	for i, row := range lds.rows {
		series[3*i] = createTimeSeries(ts, row.org, row.platform, "waiting_jobs", row.waitingJobs)
		series[3*i+1] = createTimeSeries(ts, row.org, row.platform, "running_jobs", row.runningJobs)
		series[3*i+2] = createTimeSeries(ts, row.org, row.platform, "unfinished_jobs", row.waitingJobs+row.runningJobs)
	}
	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func createTimeSeries(ts *timestamp.Timestamp, org, platform, metricType string, value int) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: fmt.Sprintf("%s/%s/%s/%s", baseMetricType, org, platform, metricType),
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(value),
				},
			},
		}},
	}
}
<commit_msg>Rename platform_load Stackdriver metric (#816)<commit_after>package metrics

import (
	"fmt"
	"time"

	"github.com/bazelbuild/continuous-integration/metrics/clients"
	"github.com/bazelbuild/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

const baseMetricType = "custom.googleapis.com/bazel/ci"

type PlatformLoad struct {
	client  clients.BuildkiteClient
	orgs    []string
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (*PlatformLoad) Type() MetricType {
	return TimeBasedMetric
}

func (*PlatformLoad) RelevantDelta() int {
	return 2 * 24 * 60 * 60 // Two days in seconds
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now(), rows: make([]*loadDataRow, 0)}
	for _, org := range pl.orgs {
		pid := &data.PipelineID{Org: org, Slug: "all"}
		builds, err := pl.client.GetMostRecentBuilds(pid, pl.builds)
		if err != nil {
			return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
		}

		allPlatforms := make(map[string]bool)
		waiting := make(map[string]int)
		running := make(map[string]int)
		for _, build := range builds {
			for _, job := range build.Jobs {
				// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
				platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
				if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
					continue
				}
				allPlatforms[platform] = true
				switch *job.State {
				case "running":
					running[platform] += 1
				case "scheduled", "runnable":
					/*
						State "scheduled" / "runnable" = waiting for a worker to become available
						State "waiting" / "waiting_failed" = waiting for another task to finish

						We're only interested in "scheduled" and "runnable" jobs since they may indicate a shortage of workers.
					*/
					waiting[platform] += 1
				}

			}
		}

		for platform := range allPlatforms {
			row := &loadDataRow{org: org, platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
			result.rows = append(result.rows, row)
		}
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, org VARCHAR(255), platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(org, timestamp, platform));
func CreatePlatformLoad(client clients.BuildkiteClient, builds int, orgs ...string) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"org", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, orgs: orgs, columns: columns, builds: builds}
}

type loadDataRow struct {
	org         string
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []*loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.org, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows)*3)
	for i, row := range lds.rows {
		series[3*i] = createTimeSeries(ts, row.org, row.platform, "waiting_jobs", row.waitingJobs)
		series[3*i+1] = createTimeSeries(ts, row.org, row.platform, "running_jobs", row.runningJobs)
		series[3*i+2] = createTimeSeries(ts, row.org, row.platform, "required_workers", row.waitingJobs+row.runningJobs)
	}
	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func createTimeSeries(ts *timestamp.Timestamp, org, platform, metricType string, value int) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: fmt.Sprintf("%s/%s/%s/%s", baseMetricType, org, platform, metricType),
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(value),
				},
			},
		}},
	}
}
<|endoftext|>