<commit_before>package client_test

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net/mail"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/client"
)

func ExampleClient() {
	log.Println("Connecting to server...")

	// Connect to server
	c, err := client.DialTLS("mail.example.org:993", nil)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Connected")

	// Don't forget to logout
	defer c.Logout()

	// Login
	if err := c.Login("username", "password"); err != nil {
		log.Fatal(err)
	}
	log.Println("Logged in")

	// List mailboxes
	mailboxes := make(chan *imap.MailboxInfo, 10)
	done := make(chan error, 1)
	go func() {
		done <- c.List("", "*", mailboxes)
	}()

	log.Println("Mailboxes:")
	for m := range mailboxes {
		log.Println("* " + m.Name)
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Flags for INBOX:", mbox.Flags)

	// Get the last 4 messages
	from := uint32(1)
	to := mbox.Messages
	if mbox.Messages > 3 {
		// We're using unsigned integers here, only substract if the result is > 0
		from = mbox.Messages - 3
	}
	seqset := new(imap.SeqSet)
	seqset.AddRange(from, to)
	items := []imap.FetchItem{imap.FetchEnvelope}

	messages := make(chan *imap.Message, 10)
	done = make(chan error, 1)
	go func() {
		done <- c.Fetch(seqset, items, messages)
	}()

	log.Println("Last 4 messages:")
	for msg := range messages {
		log.Println("* " + msg.Envelope.Subject)
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	log.Println("Done!")
}

func ExampleClient_Fetch() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// Get the last message
	if mbox.Messages == 0 {
		log.Fatal("No message in mailbox")
	}
	seqset := new(imap.SeqSet)
	seqset.AddRange(mbox.Messages, mbox.Messages)

	// Get the whole message body
	section := &imap.BodySectionName{}
	items := []imap.FetchItem{section.FetchItem()}

	messages := make(chan *imap.Message, 1)
	done := make(chan error, 1)
	go func() {
		done <- c.Fetch(seqset, items, messages)
	}()

	log.Println("Last message:")
	msg := <-messages
	r := msg.GetBody(section)
	if r == nil {
		log.Fatal("Server didn't returned message body")
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	m, err := mail.ReadMessage(r)
	if err != nil {
		log.Fatal(err)
	}

	header := m.Header
	log.Println("Date:", header.Get("Date"))
	log.Println("From:", header.Get("From"))
	log.Println("To:", header.Get("To"))
	log.Println("Subject:", header.Get("Subject"))

	body, err := ioutil.ReadAll(m.Body)
	if err != nil {
		log.Fatal(err)
	}
	log.Println(body)
}

func ExampleClient_Expunge() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// We will delete the last message
	if mbox.Messages == 0 {
		log.Fatal("No message in mailbox")
	}
	seqset := new(imap.SeqSet)
	seqset.AddNum(mbox.Messages)

	// First mark the message as deleted
	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{imap.DeletedFlag}
	if err := c.Store(seqset, item, flags, nil); err != nil {
		log.Fatal(err)
	}

	// Then delete it
	if err := c.Expunge(nil); err != nil {
		log.Fatal(err)
	}

	log.Println("Last message has been deleted")
}

func ExampleClient_StartTLS() {
	log.Println("Connecting to server...")

	// Connect to server
	c, err := client.Dial("mail.example.org:143")
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Connected")

	// Don't forget to logout
	defer c.Logout()

	// Start a TLS session
	tlsConfig := &tls.Config{ServerName: "mail.example.org"}
	if err := c.StartTLS(tlsConfig); err != nil {
		log.Fatal(err)
	}
	log.Println("TLS started")

	// Now we can login
	if err := c.Login("username", "password"); err != nil {
		log.Fatal(err)
	}
	log.Println("Logged in")
}

func ExampleClient_Store() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	_, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// Mark message 42 as seen
	seqSet := new(imap.SeqSet)
	seqSet.AddNum(42)
	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{imap.SeenFlag}
	err = c.Store(seqSet, item, flags, nil)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Message has been marked as seen")
}
<commit_msg>Add an example for Client.Search() usage<commit_after>package client_test

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net/mail"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/client"
)

func ExampleClient() {
	log.Println("Connecting to server...")

	// Connect to server
	c, err := client.DialTLS("mail.example.org:993", nil)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Connected")

	// Don't forget to logout
	defer c.Logout()

	// Login
	if err := c.Login("username", "password"); err != nil {
		log.Fatal(err)
	}
	log.Println("Logged in")

	// List mailboxes
	mailboxes := make(chan *imap.MailboxInfo, 10)
	done := make(chan error, 1)
	go func() {
		done <- c.List("", "*", mailboxes)
	}()

	log.Println("Mailboxes:")
	for m := range mailboxes {
		log.Println("* " + m.Name)
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Flags for INBOX:", mbox.Flags)

	// Get the last 4 messages
	from := uint32(1)
	to := mbox.Messages
	if mbox.Messages > 3 {
		// We're using unsigned integers here, only substract if the result is > 0
		from = mbox.Messages - 3
	}
	seqset := new(imap.SeqSet)
	seqset.AddRange(from, to)
	items := []imap.FetchItem{imap.FetchEnvelope}

	messages := make(chan *imap.Message, 10)
	done = make(chan error, 1)
	go func() {
		done <- c.Fetch(seqset, items, messages)
	}()

	log.Println("Last 4 messages:")
	for msg := range messages {
		log.Println("* " + msg.Envelope.Subject)
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	log.Println("Done!")
}

func ExampleClient_Fetch() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// Get the last message
	if mbox.Messages == 0 {
		log.Fatal("No message in mailbox")
	}
	seqset := new(imap.SeqSet)
	seqset.AddRange(mbox.Messages, mbox.Messages)

	// Get the whole message body
	section := &imap.BodySectionName{}
	items := []imap.FetchItem{section.FetchItem()}

	messages := make(chan *imap.Message, 1)
	done := make(chan error, 1)
	go func() {
		done <- c.Fetch(seqset, items, messages)
	}()

	log.Println("Last message:")
	msg := <-messages
	r := msg.GetBody(section)
	if r == nil {
		log.Fatal("Server didn't returned message body")
	}

	if err := <-done; err != nil {
		log.Fatal(err)
	}

	m, err := mail.ReadMessage(r)
	if err != nil {
		log.Fatal(err)
	}

	header := m.Header
	log.Println("Date:", header.Get("Date"))
	log.Println("From:", header.Get("From"))
	log.Println("To:", header.Get("To"))
	log.Println("Subject:", header.Get("Subject"))

	body, err := ioutil.ReadAll(m.Body)
	if err != nil {
		log.Fatal(err)
	}
	log.Println(body)
}

func ExampleClient_Expunge() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	mbox, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// We will delete the last message
	if mbox.Messages == 0 {
		log.Fatal("No message in mailbox")
	}
	seqset := new(imap.SeqSet)
	seqset.AddNum(mbox.Messages)

	// First mark the message as deleted
	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{imap.DeletedFlag}
	if err := c.Store(seqset, item, flags, nil); err != nil {
		log.Fatal(err)
	}

	// Then delete it
	if err := c.Expunge(nil); err != nil {
		log.Fatal(err)
	}

	log.Println("Last message has been deleted")
}

func ExampleClient_StartTLS() {
	log.Println("Connecting to server...")

	// Connect to server
	c, err := client.Dial("mail.example.org:143")
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Connected")

	// Don't forget to logout
	defer c.Logout()

	// Start a TLS session
	tlsConfig := &tls.Config{ServerName: "mail.example.org"}
	if err := c.StartTLS(tlsConfig); err != nil {
		log.Fatal(err)
	}
	log.Println("TLS started")

	// Now we can login
	if err := c.Login("username", "password"); err != nil {
		log.Fatal(err)
	}
	log.Println("Logged in")
}

func ExampleClient_Store() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	_, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// Mark message 42 as seen
	seqSet := new(imap.SeqSet)
	seqSet.AddNum(42)
	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{imap.SeenFlag}
	err = c.Store(seqSet, item, flags, nil)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("Message has been marked as seen")
}

func ExampleClient_Search() {
	// Let's assume c is a client
	var c *client.Client

	// Select INBOX
	_, err := c.Select("INBOX", false)
	if err != nil {
		log.Fatal(err)
	}

	// Set search criteria
	criteria := imap.NewSearchCriteria()
	criteria.WithoutFlags = []string{imap.SeenFlag}
	ids, err := c.Search(criteria)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("IDs found:", ids)

	if len(ids) > 0 {
		seqset := new(imap.SeqSet)
		seqset.AddNum(ids...)

		messages := make(chan *imap.Message, 10)
		done := make(chan error, 1)
		go func() {
			done <- c.Fetch(seqset, []imap.FetchItem{imap.FetchEnvelope}, messages)
		}()

		log.Println("Unseen messages:")
		for msg := range messages {
			log.Println("* " + msg.Envelope.Subject)
		}

		if err := <-done; err != nil {
			log.Fatal(err)
		}
	}

	log.Println("Done!")
}
<|endoftext|>


<commit_before>package client

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

	"github.com/drasko/edgex-export"
	"github.com/drasko/edgex-export/mongo"
	"github.com/go-zoo/bone"
	"go.uber.org/zap"
	"gopkg.in/mgo.v2/bson"
)

var logMsg = "Failed to query"

func getRegByID(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	id := bone.GetValue(r, "id")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := export.Registration{}
	if err := c.Find(bson.M{"id": id}).One(&reg); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg)
	if err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func getRegList(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	t := bone.GetValue(r, "type")

   var l string

	switch t {
		case "algorithms":
			l = `["None","Aes"]`
		case "compressions":
			l = `["None","Gzip","Zip"]`
		case "formats":
			l = `["JSON","XML","Serialized","IotCoreJSON","AzureJSON","CSV"]`
		case "destinations":
			l = `["DestMQTT", "TeDestZMQller", "DestIotCoreMQTT,
			"DestAzureMQTT", "DestRest"]`
		default :
			logger.Error("Unknown type: " + t);
			w.WriteHeader(http.StatusBadRequest)
			io.WriteString(w, "Unknown type: " + t)
			return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, l)
}

func getAllReg(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := []export.Registration{}
	if err := c.Find(nil).All(&reg); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg);
	if err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func getRegByName(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	name := bone.GetValue(r, "name")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := export.Registration{}
	if err := c.Find(bson.M{"name": name}).One(&reg); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg)
	if err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func addReg(w http.ResponseWriter, r *http.Request) {
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	reg := export.Registration{}
	if err := json.Unmarshal(data, &reg); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	count, err := c.Find(bson.M{"name": reg.Name}).Count();
	if  err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}
	if  count != 0 {
		logger.Error(logMsg, zap.String("Username already taken", reg.Name))
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if err := c.Insert(reg); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusCreated)
}

func updateReg(w http.ResponseWriter, r *http.Request) {
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	var body map[string]interface{}
	if err := json.Unmarshal(data, &body); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
	}

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	name := body["name"]
	query := bson.M{"name": name}
	update := bson.M{"$set": body}

	if err := c.Update(query, update); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusNotFound)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}

func delRegByID(w http.ResponseWriter, r *http.Request) {
	id := bone.GetValue(r, "id")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	if err := c.Remove(bson.M{"id": id}); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}

func delRegByName(w http.ResponseWriter, r *http.Request) {
	name := bone.GetValue(r, "name")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	if err := c.Remove(bson.M{"name": name}); err != nil {
		logger.Error(logMsg, zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}
<commit_msg>Replace logMsg by explicit in-place string<commit_after>package client

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

	"github.com/drasko/edgex-export"
	"github.com/drasko/edgex-export/mongo"
	"github.com/go-zoo/bone"
	"go.uber.org/zap"
	"gopkg.in/mgo.v2/bson"
)

func getRegByID(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	id := bone.GetValue(r, "id")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := export.Registration{}
	if err := c.Find(bson.M{"id": id}).One(&reg); err != nil {
		logger.Error("Failed to query by id", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg)
	if err != nil {
		logger.Error("Failed to query by id", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func getRegList(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	t := bone.GetValue(r, "type")

   var l string

	switch t {
		case "algorithms":
			l = `["None","Aes"]`
		case "compressions":
			l = `["None","Gzip","Zip"]`
		case "formats":
			l = `["JSON","XML","Serialized","IotCoreJSON","AzureJSON","CSV"]`
		case "destinations":
			l = `["DestMQTT", "TeDestZMQller", "DestIotCoreMQTT,
			"DestAzureMQTT", "DestRest"]`
		default :
			logger.Error("Unknown type: " + t);
			w.WriteHeader(http.StatusBadRequest)
			io.WriteString(w, "Unknown type: " + t)
			return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, l)
}

func getAllReg(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := []export.Registration{}
	if err := c.Find(nil).All(&reg); err != nil {
		logger.Error("Failed to query all registrations", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg);
	if err != nil {
		logger.Error("Failed to query all registrations", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func getRegByName(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	name := bone.GetValue(r, "name")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	reg := export.Registration{}
	if err := c.Find(bson.M{"name": name}).One(&reg); err != nil {
		logger.Error("Failed to query by name", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	res, err := json.Marshal(reg)
	if err != nil {
		logger.Error("Failed to query by name", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
	io.WriteString(w, string(res))
}

func addReg(w http.ResponseWriter, r *http.Request) {
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error("Failed to query add registration", zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	reg := export.Registration{}
	if err := json.Unmarshal(data, &reg); err != nil {
		logger.Error("Failed to query add registration", zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	count, err := c.Find(bson.M{"name": reg.Name}).Count();
	if  err != nil {
		logger.Error("Failed to query add registration", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}
	if  count != 0 {
		logger.Error("Username already taken: " + reg.Name)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if err := c.Insert(reg); err != nil {
		logger.Error("Failed to query add registration", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusCreated)
}

func updateReg(w http.ResponseWriter, r *http.Request) {
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error("Failed to query update registration", zap.Error(err))
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, err.Error())
		return
	}

	var body map[string]interface{}
	if err := json.Unmarshal(data, &body); err != nil {
		logger.Error("Failed to query update registration", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
	}

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	name := body["name"]
	query := bson.M{"name": name}
	update := bson.M{"$set": body}

	if err := c.Update(query, update); err != nil {
		logger.Error("Failed to query update registration", zap.Error(err))
		w.WriteHeader(http.StatusNotFound)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}

func delRegByID(w http.ResponseWriter, r *http.Request) {
	id := bone.GetValue(r, "id")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	if err := c.Remove(bson.M{"id": id}); err != nil {
		logger.Error("Failed to query by id", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}

func delRegByName(w http.ResponseWriter, r *http.Request) {
	name := bone.GetValue(r, "name")

	s := repo.Session.Copy()
	defer s.Close()
	c := s.DB(mongo.DbName).C(mongo.CollectionName)

	if err := c.Remove(bson.M{"name": name}); err != nil {
		logger.Error("Failed to query by name", zap.Error(err))
		w.WriteHeader(http.StatusInternalServerError)
		io.WriteString(w, err.Error())
		return
	}

	w.WriteHeader(http.StatusOK)
}
<|endoftext|>


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

package apis

import (
	"fmt"

	"github.com/apigee/apigeecli/apiclient"
	bundle "github.com/apigee/apigeecli/bundlegen"
	proxybundle "github.com/apigee/apigeecli/bundlegen/proxybundle"
	"github.com/apigee/apigeecli/client/apis"
	"github.com/spf13/cobra"
)

var OasCreateCmd = &cobra.Command{
	Use:     "openapi",
	Aliases: []string{"oas"},
	Short:   "Creates an API proxy from an OpenAPI Specification",
	Long:    "Creates an API proxy from an OpenAPI Specification",
	Args: func(cmd *cobra.Command, args []string) (err error) {
		if oasFile == "" && oasURI == "" {
			return fmt.Errorf("Either oasfile or oasuri must be passed")
		}
		return apiclient.SetApigeeOrg(org)
	},
	RunE: func(cmd *cobra.Command, args []string) (err error) {
		var content []byte
		var oasDocName string
		if oasFile != "" {
			oasDocName, content, err = bundle.LoadDocumentFromFile(oasFile, validateSpec, formatValidation)
		} else {
			oasDocName, content, err = bundle.LoadDocumentFromURI(oasURI, validateSpec, formatValidation)
		}
		if err != nil {
			return err
		}

		//Generate the apiproxy struct
		err = bundle.GenerateAPIProxyDefFromOAS(name,
			oasDocName,
			skipPolicy,
			addCORS,
			oasGoogleAcessTokenScopeLiteral,
			oasGoogleIdTokenAudLiteral,
			oasGoogleIdTokenAudRef,
			targetUrlRef,
			targetUrl)

		if err != nil {
			return err
		}

		//Create the API proxy bundle
		err = proxybundle.GenerateAPIProxyBundleFromOAS(name,
			string(content),
			oasDocName,
			skipPolicy,
			addCORS,
			oasGoogleAcessTokenScopeLiteral,
			oasGoogleIdTokenAudLiteral,
			oasGoogleIdTokenAudRef,
			targetUrlRef,
			targetUrl)

		if err != nil {
			return err
		}

		if importProxy {
			_, err = apis.CreateProxy(name, name+".zip")
		}

		return err
	},
}

var oasFile, oasURI, targetUrl string
var oasGoogleAcessTokenScopeLiteral, oasGoogleIdTokenAudLiteral, oasGoogleIdTokenAudRef string
var validateSpec, formatValidation bool

func init() {
	OasCreateCmd.Flags().StringVarP(&name, "name", "n",
		"", "API Proxy name")
	OasCreateCmd.Flags().StringVarP(&oasFile, "oasfile", "f",
		"", "Open API 3.0 Specification file")
	OasCreateCmd.Flags().StringVarP(&oasURI, "oasuri", "u",
		"", "Open API 3.0 Specification URI location")
	OasCreateCmd.Flags().StringVarP(&oasGoogleAcessTokenScopeLiteral, "google-accesstoken-scope-literal", "",
		"", "Generate Google Access token with target endpoint and set scope")
	OasCreateCmd.Flags().StringVarP(&oasGoogleIdTokenAudLiteral, "google-idtoken-aud-literal", "",
		"", "Generate Google ID token with target endpoint and set audience")
	OasCreateCmd.Flags().StringVarP(&oasGoogleIdTokenAudRef, "google-idtoken-aud-ref", "",
		"", "Generate Google ID token token with target endpoint and set audience reference")
	OasCreateCmd.Flags().StringVarP(&targetUrlRef, "target-url-ref", "",
		"", "Set a reference variable containing the target endpoint")
	OasCreateCmd.Flags().StringVarP(&targetUrl, "target-url", "",
		"", "Set a target URL for the target endpoint")
	OasCreateCmd.Flags().BoolVarP(&importProxy, "import", "",
		true, "Import API Proxy after generation from spec")
	OasCreateCmd.Flags().BoolVarP(&validateSpec, "validate", "",
		true, "Validate Spec before generating proxy")
	OasCreateCmd.Flags().BoolVarP(&skipPolicy, "skip-policy", "",
		false, "Skip adding the OAS Validate policy")
	OasCreateCmd.Flags().BoolVarP(&addCORS, "add-cors", "",
		false, "Add a CORS policy")
	OasCreateCmd.Flags().BoolVarP(&formatValidation, "formatValidation", "",
		true, "disables validation of schema type formats")

	_ = OasCreateCmd.MarkFlagRequired("name")
}
<commit_msg>validations for issue #20<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package apis

import (
	"fmt"

	"github.com/apigee/apigeecli/apiclient"
	bundle "github.com/apigee/apigeecli/bundlegen"
	proxybundle "github.com/apigee/apigeecli/bundlegen/proxybundle"
	"github.com/apigee/apigeecli/client/apis"
	"github.com/spf13/cobra"
)

var OasCreateCmd = &cobra.Command{
	Use:     "openapi",
	Aliases: []string{"oas"},
	Short:   "Creates an API proxy from an OpenAPI Specification",
	Long:    "Creates an API proxy from an OpenAPI Specification",
	Args: func(cmd *cobra.Command, args []string) (err error) {
		if oasFile == "" && oasURI == "" {
			return fmt.Errorf("Either oasfile or oasuri must be passed")
		}
		if targetUrl != "" && targetUrlRef != "" {
			return fmt.Errorf("Either target-url or target-url-ref must be passed, not both")
		}
		return apiclient.SetApigeeOrg(org)
	},
	RunE: func(cmd *cobra.Command, args []string) (err error) {
		var content []byte
		var oasDocName string
		if oasFile != "" {
			oasDocName, content, err = bundle.LoadDocumentFromFile(oasFile, validateSpec, formatValidation)
		} else {
			oasDocName, content, err = bundle.LoadDocumentFromURI(oasURI, validateSpec, formatValidation)
		}
		if err != nil {
			return err
		}

		//Generate the apiproxy struct
		err = bundle.GenerateAPIProxyDefFromOAS(name,
			oasDocName,
			skipPolicy,
			addCORS,
			oasGoogleAcessTokenScopeLiteral,
			oasGoogleIdTokenAudLiteral,
			oasGoogleIdTokenAudRef,
			targetUrlRef,
			targetUrl)

		if err != nil {
			return err
		}

		//Create the API proxy bundle
		err = proxybundle.GenerateAPIProxyBundleFromOAS(name,
			string(content),
			oasDocName,
			skipPolicy,
			addCORS,
			oasGoogleAcessTokenScopeLiteral,
			oasGoogleIdTokenAudLiteral,
			oasGoogleIdTokenAudRef,
			targetUrlRef,
			targetUrl)

		if err != nil {
			return err
		}

		if importProxy {
			_, err = apis.CreateProxy(name, name+".zip")
		}

		return err
	},
}

var oasFile, oasURI, targetUrl string
var oasGoogleAcessTokenScopeLiteral, oasGoogleIdTokenAudLiteral, oasGoogleIdTokenAudRef string
var validateSpec, formatValidation bool

func init() {
	OasCreateCmd.Flags().StringVarP(&name, "name", "n",
		"", "API Proxy name")
	OasCreateCmd.Flags().StringVarP(&oasFile, "oasfile", "f",
		"", "Open API 3.0 Specification file")
	OasCreateCmd.Flags().StringVarP(&oasURI, "oasuri", "u",
		"", "Open API 3.0 Specification URI location")
	OasCreateCmd.Flags().StringVarP(&oasGoogleAcessTokenScopeLiteral, "google-accesstoken-scope-literal", "",
		"", "Generate Google Access token with target endpoint and set scope")
	OasCreateCmd.Flags().StringVarP(&oasGoogleIdTokenAudLiteral, "google-idtoken-aud-literal", "",
		"", "Generate Google ID token with target endpoint and set audience")
	OasCreateCmd.Flags().StringVarP(&oasGoogleIdTokenAudRef, "google-idtoken-aud-ref", "",
		"", "Generate Google ID token token with target endpoint and set audience reference")
	OasCreateCmd.Flags().StringVarP(&targetUrlRef, "target-url-ref", "",
		"", "Set a reference variable containing the target endpoint")
	OasCreateCmd.Flags().StringVarP(&targetUrl, "target-url", "",
		"", "Set a target URL for the target endpoint")
	OasCreateCmd.Flags().BoolVarP(&importProxy, "import", "",
		true, "Import API Proxy after generation from spec")
	OasCreateCmd.Flags().BoolVarP(&validateSpec, "validate", "",
		true, "Validate Spec before generating proxy")
	OasCreateCmd.Flags().BoolVarP(&skipPolicy, "skip-policy", "",
		false, "Skip adding the OAS Validate policy")
	OasCreateCmd.Flags().BoolVarP(&addCORS, "add-cors", "",
		false, "Add a CORS policy")
	OasCreateCmd.Flags().BoolVarP(&formatValidation, "formatValidation", "",
		true, "disables validation of schema type formats")

	_ = OasCreateCmd.MarkFlagRequired("name")
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"

	humanize "github.com/dustin/go-humanize"
	"github.com/spf13/cobra"
	"github.com/stormforger/cli/api/filefixture"
)

var (
	datasourceShowCmd = &cobra.Command{
		Use:     "show <organisation-ref> <name>",
		Aliases: []string{},
		Short:   "Show details of fixture",
		Run:     runDatasourceShow,
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if len(args) > 2 {
				log.Fatal("Too many arguments")
			}

			if len(args) < 2 {
				log.Fatal("Missing organisation or datasource")
			}

			datasourceOpts.Organisation = lookupOrganisationUID(*NewClient(), args[0])
			if datasourceOpts.Organisation == "" {
				log.Fatal("Missing organisation")
			}
		},
	}
)

func init() {
	datasourceCmd.AddCommand(datasourceShowCmd)
}

func runDatasourceShow(cmd *cobra.Command, args []string) {
	client := NewClient()
	fileName := args[1]

	fileFixture := findFixtureByName(*client, datasourceOpts.Organisation, fileName)

	ShowDetails(fileFixture)
}

// ShowDetails print out details of a file fixture, including its current version
func ShowDetails(fileFixture *filefixture.FileFixture) {
	fmt.Printf("Name:            %s\n", fileFixture.Name)
	fmt.Printf("UID:             %s\n", fileFixture.ID)
	fmt.Printf("Created:         %s (%s)\n", convertToLocalTZ(fileFixture.CreatedAt), humanize.Time(fileFixture.CreatedAt))
	fmt.Printf("Updated:         %s (%s)\n", convertToLocalTZ(fileFixture.UpdatedAt), humanize.Time(fileFixture.UpdatedAt))
	fmt.Printf("Current Version: %s\n", fileFixture.CurrentVersion.ID)
	fmt.Printf("  SHA256 Hash:   %s\n", fileFixture.CurrentVersion.Hash)
	fmt.Printf("  Size:          %s\n", humanize.Bytes(uint64(fileFixture.CurrentVersion.FileSize)))
	fmt.Printf("  Line Count:    %v\n", fileFixture.CurrentVersion.ItemCount)
	fmt.Printf("  Created:       %s (%s)\n", convertToLocalTZ(fileFixture.CurrentVersion.CreatedAt), humanize.Time(fileFixture.CurrentVersion.CreatedAt))
	fmt.Printf("Version(s):      %v\n", len(fileFixture.Versions))
	for _, version := range fileFixture.Versions {
		fmt.Printf("  - %s (created %s)\n", version.ID, convertToLocalTZ(version.CreatedAt))
	}
}
<commit_msg>refactor: simplify arg handling<commit_after>package cmd

import (
	"fmt"
	"log"

	humanize "github.com/dustin/go-humanize"
	"github.com/spf13/cobra"
	"github.com/stormforger/cli/api/filefixture"
)

var (
	datasourceShowCmd = &cobra.Command{
		Use:     "show <organisation-ref> <name>",
		Aliases: []string{},
		Short:   "Show details of fixture",
		Args:    cobra.ExactArgs(2),
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			datasourceOpts.Organisation = lookupOrganisationUID(*NewClient(), args[0])
			if datasourceOpts.Organisation == "" {
				log.Fatal("Missing organisation")
			}
		},
		Run: runDatasourceShow,
	}
)

func init() {
	datasourceCmd.AddCommand(datasourceShowCmd)
}

func runDatasourceShow(cmd *cobra.Command, args []string) {
	client := NewClient()
	fileName := args[1]

	fileFixture := findFixtureByName(*client, datasourceOpts.Organisation, fileName)

	ShowDetails(fileFixture)
}

// ShowDetails print out details of a file fixture, including its current version
func ShowDetails(fileFixture *filefixture.FileFixture) {
	fmt.Printf("Name:            %s\n", fileFixture.Name)
	fmt.Printf("UID:             %s\n", fileFixture.ID)
	fmt.Printf("Created:         %s (%s)\n", convertToLocalTZ(fileFixture.CreatedAt), humanize.Time(fileFixture.CreatedAt))
	fmt.Printf("Updated:         %s (%s)\n", convertToLocalTZ(fileFixture.UpdatedAt), humanize.Time(fileFixture.UpdatedAt))
	fmt.Printf("Current Version: %s\n", fileFixture.CurrentVersion.ID)
	fmt.Printf("  SHA256 Hash:   %s\n", fileFixture.CurrentVersion.Hash)
	fmt.Printf("  Size:          %s\n", humanize.Bytes(uint64(fileFixture.CurrentVersion.FileSize)))
	fmt.Printf("  Line Count:    %v\n", fileFixture.CurrentVersion.ItemCount)
	fmt.Printf("  Created:       %s (%s)\n", convertToLocalTZ(fileFixture.CurrentVersion.CreatedAt), humanize.Time(fileFixture.CurrentVersion.CreatedAt))
	fmt.Printf("Version(s):      %v\n", len(fileFixture.Versions))
	for _, version := range fileFixture.Versions {
		fmt.Printf("  - %s (created %s)\n", version.ID, convertToLocalTZ(version.CreatedAt))
	}
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/signal"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui  cli.Ui
	Cmd string

	offset   int64
	lastN    int64 // peek the most recent N messages
	colorize bool
	limit    int
	quit     chan struct{}
	once     sync.Once
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster      string
		zone         string
		topicPattern string
		partitionId  int
		silence      bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topicPattern, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&this.colorize, "color", true, "")
	cmdFlags.Int64Var(&this.lastN, "last", -1, "")
	cmdFlags.IntVar(&this.limit, "limit", -1, "")
	cmdFlags.Int64Var(&this.offset, "offset", sarama.OffsetNewest, "")
	cmdFlags.BoolVar(&silence, "s", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.quit = make(chan struct{})

	if silence {
		stats := newPeekStats()
		go stats.start()
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 20000) // msg aggerator channel
	if cluster == "" {
		zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
			this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
	}

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Printf("received signal: %s", strings.ToUpper(sig.String()))
		log.Println("quiting...")

		this.once.Do(func() {
			close(this.quit)
		})
	}, syscall.SIGINT, syscall.SIGTERM)

	var (
		startAt = time.Now()
		msg     *sarama.ConsumerMessage
		total   int
		bytes   int64
	)

LOOP:
	for {
		select {
		case <-this.quit:
			this.Ui.Output(fmt.Sprintf("Total: %s msgs, %s", gofmt.Comma(int64(total)), gofmt.ByteSize(bytes)))
			this.Ui.Output(fmt.Sprintf("Speed: %d/s", total/int(time.Since(startAt).Seconds())))
			return

		case msg = <-msgChan:
			if silence {
				stats.MsgCountPerSecond.Mark(1)
				stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))
			} else {
				if this.colorize {
					this.Ui.Output(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						color.Green(msg.Topic), msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				} else {
					// colored UI will have invisible chars output
					fmt.Println(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						msg.Topic, msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				}
			}

			total++
			bytes += int64(len(msg.Value))

			if this.limit > 0 && total >= this.limit {
				break LOOP

			}
			if this.lastN > 0 && total >= int(this.lastN) {
				break LOOP
			}

		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topicPattern string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	topics, err := kfk.Topics()
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}

	for _, t := range topics {
		if patternMatched(t, topicPattern) {
			go this.simpleConsumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	}

}

func (this *Peek) simpleConsumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			offset := this.offset
			if this.lastN > 0 {
				latestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetNewest)
				swallow(err)

				oldestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetOldest)
				swallow(err)

				offset = latestOffset - this.lastN
				if offset < oldestOffset {
					offset = oldestOffset
				}

				if offset == 0 {
					// no message in store
					return
				}
			}

			go this.consumePartition(kfk, consumer, topic, p, msgCh, offset)
		}

	} else {
		offset := this.offset
		if this.lastN > 0 {
			latestOffset, err := kfk.GetOffset(topic, partitionId, sarama.OffsetNewest)
			swallow(err)
			offset = latestOffset - this.lastN
			if offset < 0 {
				offset = sarama.OffsetOldest
			}
		}
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh, offset)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage, offset int64) {
	p, err := consumer.ConsumePartition(topic, partitionId, offset)
	if err != nil {
		this.Ui.Error(fmt.Sprintf("%s/%d: offset=%d %v", topic, partitionId, offset, err))
		return
	}
	defer p.Close()

	n := int64(0)
	for {
		select {
		case <-this.quit:
			return

		case msg := <-p.Messages():
			msgCh <- msg

			n++
			if this.lastN > 0 && n >= this.lastN {
				return
			}
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing from any offset"
}

func (this *Peek) Help() string {
	help := fmt.Sprintf(`
Usage: %s peek [options]

    Peek kafka cluster messages ongoing from any offset

Options:

    -z zone
      Default %s

    -c cluster

    -t topic pattern

    -p partition id
      -1 will peek all partitions of a topic

    -last n
      Peek the most recent N messages

    -offset message offset value
      -1 OffsetNewest, -2 OffsetOldest.
      You can specify your own offset.
      Default -1(OffsetNewest)

    -limit n
      Limit how many messages to consume

    -s
      Silence mode, only display statastics instead of message content

    -color
      Enable colorized output
`, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<commit_msg>avoid divide by zero panic<commit_after>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/signal"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui  cli.Ui
	Cmd string

	offset   int64
	lastN    int64 // peek the most recent N messages
	colorize bool
	limit    int
	quit     chan struct{}
	once     sync.Once
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster      string
		zone         string
		topicPattern string
		partitionId  int
		silence      bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topicPattern, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&this.colorize, "color", true, "")
	cmdFlags.Int64Var(&this.lastN, "last", -1, "")
	cmdFlags.IntVar(&this.limit, "limit", -1, "")
	cmdFlags.Int64Var(&this.offset, "offset", sarama.OffsetNewest, "")
	cmdFlags.BoolVar(&silence, "s", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.quit = make(chan struct{})

	if silence {
		stats := newPeekStats()
		go stats.start()
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 20000) // msg aggerator channel
	if cluster == "" {
		zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
			this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
	}

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Printf("received signal: %s", strings.ToUpper(sig.String()))
		log.Println("quiting...")

		this.once.Do(func() {
			close(this.quit)
		})
	}, syscall.SIGINT, syscall.SIGTERM)

	var (
		startAt = time.Now()
		msg     *sarama.ConsumerMessage
		total   int
		bytes   int64
	)

LOOP:
	for {
		select {
		case <-this.quit:
			this.Ui.Output(fmt.Sprintf("Total: %s msgs, %s", gofmt.Comma(int64(total)), gofmt.ByteSize(bytes)))
			elapsed := time.Since(startAt).Seconds()
			if elapsed > 0. {
				this.Ui.Output(fmt.Sprintf("Speed: %d/s", total/int(elapsed)))
			}

			return

		case msg = <-msgChan:
			if silence {
				stats.MsgCountPerSecond.Mark(1)
				stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))
			} else {
				if this.colorize {
					this.Ui.Output(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						color.Green(msg.Topic), msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				} else {
					// colored UI will have invisible chars output
					fmt.Println(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						msg.Topic, msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				}
			}

			total++
			bytes += int64(len(msg.Value))

			if this.limit > 0 && total >= this.limit {
				break LOOP

			}
			if this.lastN > 0 && total >= int(this.lastN) {
				break LOOP
			}

		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topicPattern string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	topics, err := kfk.Topics()
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}

	for _, t := range topics {
		if patternMatched(t, topicPattern) {
			go this.simpleConsumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	}

}

func (this *Peek) simpleConsumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			offset := this.offset
			if this.lastN > 0 {
				latestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetNewest)
				swallow(err)

				oldestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetOldest)
				swallow(err)

				offset = latestOffset - this.lastN
				if offset < oldestOffset {
					offset = oldestOffset
				}

				if offset == 0 {
					// no message in store
					return
				}
			}

			go this.consumePartition(kfk, consumer, topic, p, msgCh, offset)
		}

	} else {
		offset := this.offset
		if this.lastN > 0 {
			latestOffset, err := kfk.GetOffset(topic, partitionId, sarama.OffsetNewest)
			swallow(err)
			offset = latestOffset - this.lastN
			if offset < 0 {
				offset = sarama.OffsetOldest
			}
		}
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh, offset)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage, offset int64) {
	p, err := consumer.ConsumePartition(topic, partitionId, offset)
	if err != nil {
		this.Ui.Error(fmt.Sprintf("%s/%d: offset=%d %v", topic, partitionId, offset, err))
		return
	}
	defer p.Close()

	n := int64(0)
	for {
		select {
		case <-this.quit:
			return

		case msg := <-p.Messages():
			msgCh <- msg

			n++
			if this.lastN > 0 && n >= this.lastN {
				return
			}
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing from any offset"
}

func (this *Peek) Help() string {
	help := fmt.Sprintf(`
Usage: %s peek [options]

    Peek kafka cluster messages ongoing from any offset

Options:

    -z zone
      Default %s

    -c cluster

    -t topic pattern

    -p partition id
      -1 will peek all partitions of a topic

    -last n
      Peek the most recent N messages

    -offset message offset value
      -1 OffsetNewest, -2 OffsetOldest.
      You can specify your own offset.
      Default -1(OffsetNewest)

    -limit n
      Limit how many messages to consume

    -s
      Silence mode, only display statastics instead of message content

    -color
      Enable colorized output
`, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kube_client "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/kubecfg"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/version"
	verflag "github.com/GoogleCloudPlatform/kubernetes/pkg/version/flag"
	"github.com/golang/glog"
)

var (
	serverVersion = flag.Bool("server_version", false, "Print the server's version number.")
	preventSkew   = flag.Bool("expect_version_match", false, "Fail if server's version doesn't match own version.")
	httpServer    = flag.String("h", "", "The host to connect to.")
	config        = flag.String("c", "", "Path to the config file.")
	selector      = flag.String("l", "", "Selector (label query) to use for listing")
	updatePeriod  = flag.Duration("u", 60*time.Second, "Update interval period")
	portSpec      = flag.String("p", "", "The port spec, comma-separated list of <external>:<internal>,...")
	servicePort   = flag.Int("s", -1, "If positive, create and run a corresponding service on this port, only used with 'run'")
	authConfig    = flag.String("auth", os.Getenv("HOME")+"/.kubernetes_auth", "Path to the auth info file.  If missing, prompt the user.  Only used if doing https.")
	json          = flag.Bool("json", false, "If true, print raw JSON for responses")
	yaml          = flag.Bool("yaml", false, "If true, print raw YAML for responses")
	verbose       = flag.Bool("verbose", false, "If true, print extra information")
	proxy         = flag.Bool("proxy", false, "If true, run a proxy to the api server")
	www           = flag.String("www", "", "If -proxy is true, use this directory to serve static files")
	templateFile  = flag.String("template_file", "", "If present, load this file as a golang template and use it for output printing")
	templateStr   = flag.String("template", "", "If present, parse this string as a golang template and use it for output printing")
)

var parser = kubecfg.NewParser(map[string]interface{}{
	"pods":                   api.Pod{},
	"services":               api.Service{},
	"replicationControllers": api.ReplicationController{},
	"minions":                api.Minion{},
})

func usage() {
	fmt.Fprintf(os.Stderr, `usage: kubecfg -h [-c config/file.json] [-p :,..., :] <method>

  Kubernetes REST API:
  kubecfg [OPTIONS] get|list|create|delete|update <%s>[/<id>]

  Manage replication controllers:
  kubecfg [OPTIONS] stop|rm|rollingupdate <controller>
  kubecfg [OPTIONS] run <image> <replicas> <controller>
  kubecfg [OPTIONS] resize <controller> <replicas>

  Options:
`, prettyWireStorage())
	flag.PrintDefaults()

}

func prettyWireStorage() string {
	types := parser.SupportedWireStorage()
	sort.Strings(types)
	return strings.Join(types, "|")
}

// readConfig reads and parses pod, replicationController, and service
// configuration files. If any errors log and exit non-zero.
func readConfig(storage string) []byte {
	if len(*config) == 0 {
		glog.Fatal("Need config file (-c)")
	}
	data, err := ioutil.ReadFile(*config)
	if err != nil {
		glog.Fatalf("Unable to read %v: %v\n", *config, err)
	}
	data, err = parser.ToWireFormat(data, storage)
	if err != nil {
		glog.Fatalf("Error parsing %v as an object for %v: %v\n", *config, storage, err)
	}
	if *verbose {
		glog.Infof("Parsed config file successfully; sending:\n%v\n", string(data))
	}
	return data
}

func main() {
	flag.Usage = func() {
		usage()
	}

	flag.Parse()
	util.InitLogs()
	defer util.FlushLogs()

	verflag.PrintAndExitIfRequested()

	secure := true
	var masterServer string
	if len(*httpServer) > 0 {
		masterServer = *httpServer
	} else if len(os.Getenv("KUBERNETES_MASTER")) > 0 {
		masterServer = os.Getenv("KUBERNETES_MASTER")
	} else {
		masterServer = "http://localhost:8080"
	}
	parsedURL, err := url.Parse(masterServer)
	if err != nil {
		glog.Fatalf("Unable to parse %v as a URL\n", err)
	}
	if parsedURL.Scheme != "" && parsedURL.Scheme != "https" {
		secure = false
	}

	var auth *kube_client.AuthInfo
	if secure {
		auth, err = kubecfg.LoadAuthInfo(*authConfig, os.Stdin)
		if err != nil {
			glog.Fatalf("Error loading auth: %v", err)
		}
	}

	client := kube_client.New(masterServer, auth)

	if *serverVersion {
		got, err := client.ServerVersion()
		if err != nil {
			fmt.Printf("Couldn't read version from server: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("Server Version: %#v\n", got)
		os.Exit(0)
	}

	if *preventSkew {
		got, err := client.ServerVersion()
		if err != nil {
			fmt.Printf("Couldn't read version from server: %v\n", err)
			os.Exit(1)
		}
		if c, s := version.Get(), *got; !reflect.DeepEqual(c, s) {
			fmt.Printf("Server version (%#v) differs from client version (%#v)!\n", s, c)
			os.Exit(1)
		}
	}

	if *proxy {
		glog.Info("Starting to serve on localhost:8001")
		server := kubecfg.NewProxyServer(*www, masterServer, auth)
		glog.Fatal(server.Serve())
	}

	if len(flag.Args()) < 1 {
		usage()
		os.Exit(1)
	}
	method := flag.Arg(0)

	matchFound := executeAPIRequest(method, client) || executeControllerRequest(method, client)
	if matchFound == false {
		glog.Fatalf("Unknown command %s", method)
	}
}

// storagePathFromArg normalizes a path and breaks out the first segment if available
func storagePathFromArg(arg string) (storage, path string, hasSuffix bool) {
	path = strings.Trim(arg, "/")
	segments := strings.SplitN(path, "/", 2)
	storage = segments[0]
	if len(segments) > 1 && segments[1] != "" {
		hasSuffix = true
	}
	return storage, path, hasSuffix
}

//checkStorage returns true if the provided storage is valid
func checkStorage(storage string) bool {
	for _, allowed := range parser.SupportedWireStorage() {
		if allowed == storage {
			return true
		}
	}
	return false
}

func executeAPIRequest(method string, s *kube_client.Client) bool {
	storage, path, hasSuffix := storagePathFromArg(flag.Arg(1))
	validStorage := checkStorage(storage)
	verb := ""
	setBody := false
	var version uint64
	switch method {
	case "get":
		verb = "GET"
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>[/<id>]", method, prettyWireStorage())
		}
	case "list":
		verb = "GET"
		if !validStorage || hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>", method, prettyWireStorage())
		}
	case "delete":
		verb = "DELETE"
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>/<id>", method, prettyWireStorage())
		}
	case "create":
		verb = "POST"
		setBody = true
		if !validStorage || hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>", method, prettyWireStorage())
		}
	case "update":
		obj, err := s.Verb("GET").Path(path).Do().Get()
		if err != nil {
			glog.Fatalf("error obtaining resource version for update: %v", err)
		}
		jsonBase, err := api.FindJSONBase(obj)
		if err != nil {
			glog.Fatalf("error finding json base for update: %v", err)
		}
		version = jsonBase.ResourceVersion()
		verb = "PUT"
		setBody = true
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>/<id>", method, prettyWireStorage())
		}
	default:
		return false
	}

	r := s.Verb(verb).
		Path(path).
		ParseSelectorParam("labels", *selector)
	if setBody {
		if version != 0 {
			data := readConfig(storage)
			obj, err := api.Decode(data)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			jsonBase, err := api.FindJSONBase(obj)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			jsonBase.SetResourceVersion(version)
			data, err = api.Encode(obj)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			r.Body(data)
		} else {
			r.Body(readConfig(storage))
		}
	}
	result := r.Do()
	obj, err := result.Get()
	if err != nil {
		glog.Fatalf("Got request error: %v\n", err)
		return false
	}

	var printer kubecfg.ResourcePrinter
	switch {
	case *json:
		printer = &kubecfg.IdentityPrinter{}
	case *yaml:
		printer = &kubecfg.YAMLPrinter{}
	case len(*templateFile) > 0 || len(*templateStr) > 0:
		var data []byte
		if len(*templateFile) > 0 {
			var err error
			data, err = ioutil.ReadFile(*templateFile)
			if err != nil {
				glog.Fatalf("Error reading template %s, %v\n", *templateFile, err)
				return false
			}
		} else {
			data = []byte(*templateStr)
		}
		tmpl, err := template.New("output").Parse(string(data))
		if err != nil {
			glog.Fatalf("Error parsing template %s, %v\n", string(data), err)
			return false
		}
		printer = &kubecfg.TemplatePrinter{
			Template: tmpl,
		}
	default:
		printer = humanReadablePrinter()
	}

	if err = printer.PrintObj(obj, os.Stdout); err != nil {
		body, _ := result.Raw()
		glog.Fatalf("Failed to print: %v\nRaw received object:\n%#v\n\nBody received: %v", err, obj, string(body))
	}
	fmt.Print("\n")

	return true
}

func executeControllerRequest(method string, c *kube_client.Client) bool {
	parseController := func() string {
		if len(flag.Args()) != 2 {
			glog.Fatal("usage: kubecfg [OPTIONS] stop|rm|rollingupdate <controller>")
		}
		return flag.Arg(1)
	}

	var err error
	switch method {
	case "stop":
		err = kubecfg.StopController(parseController(), c)
	case "rm":
		err = kubecfg.DeleteController(parseController(), c)
	case "rollingupdate":
		err = kubecfg.Update(parseController(), c, *updatePeriod)
	case "run":
		if len(flag.Args()) != 4 {
			glog.Fatal("usage: kubecfg [OPTIONS] run <image> <replicas> <controller>")
		}
		image := flag.Arg(1)
		replicas, err := strconv.Atoi(flag.Arg(2))
		name := flag.Arg(3)
		if err != nil {
			glog.Fatalf("Error parsing replicas: %v", err)
		}
		err = kubecfg.RunController(image, name, replicas, c, *portSpec, *servicePort)
	case "resize":
		args := flag.Args()
		if len(args) < 3 {
			glog.Fatal("usage: kubecfg resize <controller> <replicas>")
		}
		name := args[1]
		replicas, err := strconv.Atoi(args[2])
		if err != nil {
			glog.Fatalf("Error parsing replicas: %v", err)
		}
		err = kubecfg.ResizeController(name, replicas, c)
	default:
		return false
	}
	if err != nil {
		glog.Fatalf("Error: %v", err)
	}
	return true
}

func humanReadablePrinter() *kubecfg.HumanReadablePrinter {
	printer := kubecfg.NewHumanReadablePrinter()
	// Add Handler calls here to support additional types
	return printer
}
<commit_msg>Don't shadow important variables.<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kube_client "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/kubecfg"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/version"
	verflag "github.com/GoogleCloudPlatform/kubernetes/pkg/version/flag"
	"github.com/golang/glog"
)

var (
	serverVersion = flag.Bool("server_version", false, "Print the server's version number.")
	preventSkew   = flag.Bool("expect_version_match", false, "Fail if server's version doesn't match own version.")
	httpServer    = flag.String("h", "", "The host to connect to.")
	config        = flag.String("c", "", "Path to the config file.")
	selector      = flag.String("l", "", "Selector (label query) to use for listing")
	updatePeriod  = flag.Duration("u", 60*time.Second, "Update interval period")
	portSpec      = flag.String("p", "", "The port spec, comma-separated list of <external>:<internal>,...")
	servicePort   = flag.Int("s", -1, "If positive, create and run a corresponding service on this port, only used with 'run'")
	authConfig    = flag.String("auth", os.Getenv("HOME")+"/.kubernetes_auth", "Path to the auth info file.  If missing, prompt the user.  Only used if doing https.")
	json          = flag.Bool("json", false, "If true, print raw JSON for responses")
	yaml          = flag.Bool("yaml", false, "If true, print raw YAML for responses")
	verbose       = flag.Bool("verbose", false, "If true, print extra information")
	proxy         = flag.Bool("proxy", false, "If true, run a proxy to the api server")
	www           = flag.String("www", "", "If -proxy is true, use this directory to serve static files")
	templateFile  = flag.String("template_file", "", "If present, load this file as a golang template and use it for output printing")
	templateStr   = flag.String("template", "", "If present, parse this string as a golang template and use it for output printing")
)

var parser = kubecfg.NewParser(map[string]interface{}{
	"pods":                   api.Pod{},
	"services":               api.Service{},
	"replicationControllers": api.ReplicationController{},
	"minions":                api.Minion{},
})

func usage() {
	fmt.Fprintf(os.Stderr, `usage: kubecfg -h [-c config/file.json] [-p :,..., :] <method>

  Kubernetes REST API:
  kubecfg [OPTIONS] get|list|create|delete|update <%s>[/<id>]

  Manage replication controllers:
  kubecfg [OPTIONS] stop|rm|rollingupdate <controller>
  kubecfg [OPTIONS] run <image> <replicas> <controller>
  kubecfg [OPTIONS] resize <controller> <replicas>

  Options:
`, prettyWireStorage())
	flag.PrintDefaults()

}

func prettyWireStorage() string {
	types := parser.SupportedWireStorage()
	sort.Strings(types)
	return strings.Join(types, "|")
}

// readConfig reads and parses pod, replicationController, and service
// configuration files. If any errors log and exit non-zero.
func readConfig(storage string) []byte {
	if len(*config) == 0 {
		glog.Fatal("Need config file (-c)")
	}
	data, err := ioutil.ReadFile(*config)
	if err != nil {
		glog.Fatalf("Unable to read %v: %v\n", *config, err)
	}
	data, err = parser.ToWireFormat(data, storage)
	if err != nil {
		glog.Fatalf("Error parsing %v as an object for %v: %v\n", *config, storage, err)
	}
	if *verbose {
		glog.Infof("Parsed config file successfully; sending:\n%v\n", string(data))
	}
	return data
}

func main() {
	flag.Usage = func() {
		usage()
	}

	flag.Parse()
	util.InitLogs()
	defer util.FlushLogs()

	verflag.PrintAndExitIfRequested()

	secure := true
	var masterServer string
	if len(*httpServer) > 0 {
		masterServer = *httpServer
	} else if len(os.Getenv("KUBERNETES_MASTER")) > 0 {
		masterServer = os.Getenv("KUBERNETES_MASTER")
	} else {
		masterServer = "http://localhost:8080"
	}
	parsedURL, err := url.Parse(masterServer)
	if err != nil {
		glog.Fatalf("Unable to parse %v as a URL\n", err)
	}
	if parsedURL.Scheme != "" && parsedURL.Scheme != "https" {
		secure = false
	}

	var auth *kube_client.AuthInfo
	if secure {
		auth, err = kubecfg.LoadAuthInfo(*authConfig, os.Stdin)
		if err != nil {
			glog.Fatalf("Error loading auth: %v", err)
		}
	}

	client := kube_client.New(masterServer, auth)

	if *serverVersion {
		got, err := client.ServerVersion()
		if err != nil {
			fmt.Printf("Couldn't read version from server: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("Server Version: %#v\n", got)
		os.Exit(0)
	}

	if *preventSkew {
		got, err := client.ServerVersion()
		if err != nil {
			fmt.Printf("Couldn't read version from server: %v\n", err)
			os.Exit(1)
		}
		if c, s := version.Get(), *got; !reflect.DeepEqual(c, s) {
			fmt.Printf("Server version (%#v) differs from client version (%#v)!\n", s, c)
			os.Exit(1)
		}
	}

	if *proxy {
		glog.Info("Starting to serve on localhost:8001")
		server := kubecfg.NewProxyServer(*www, masterServer, auth)
		glog.Fatal(server.Serve())
	}

	if len(flag.Args()) < 1 {
		usage()
		os.Exit(1)
	}
	method := flag.Arg(0)

	matchFound := executeAPIRequest(method, client) || executeControllerRequest(method, client)
	if matchFound == false {
		glog.Fatalf("Unknown command %s", method)
	}
}

// storagePathFromArg normalizes a path and breaks out the first segment if available
func storagePathFromArg(arg string) (storage, path string, hasSuffix bool) {
	path = strings.Trim(arg, "/")
	segments := strings.SplitN(path, "/", 2)
	storage = segments[0]
	if len(segments) > 1 && segments[1] != "" {
		hasSuffix = true
	}
	return storage, path, hasSuffix
}

//checkStorage returns true if the provided storage is valid
func checkStorage(storage string) bool {
	for _, allowed := range parser.SupportedWireStorage() {
		if allowed == storage {
			return true
		}
	}
	return false
}

func executeAPIRequest(method string, s *kube_client.Client) bool {
	storage, path, hasSuffix := storagePathFromArg(flag.Arg(1))
	validStorage := checkStorage(storage)
	verb := ""
	setBody := false
	var version uint64
	switch method {
	case "get":
		verb = "GET"
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>[/<id>]", method, prettyWireStorage())
		}
	case "list":
		verb = "GET"
		if !validStorage || hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>", method, prettyWireStorage())
		}
	case "delete":
		verb = "DELETE"
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>/<id>", method, prettyWireStorage())
		}
	case "create":
		verb = "POST"
		setBody = true
		if !validStorage || hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>", method, prettyWireStorage())
		}
	case "update":
		obj, err := s.Verb("GET").Path(path).Do().Get()
		if err != nil {
			glog.Fatalf("error obtaining resource version for update: %v", err)
		}
		jsonBase, err := api.FindJSONBase(obj)
		if err != nil {
			glog.Fatalf("error finding json base for update: %v", err)
		}
		version = jsonBase.ResourceVersion()
		verb = "PUT"
		setBody = true
		if !validStorage || !hasSuffix {
			glog.Fatalf("usage: kubecfg [OPTIONS] %s <%s>/<id>", method, prettyWireStorage())
		}
	default:
		return false
	}

	r := s.Verb(verb).
		Path(path).
		ParseSelectorParam("labels", *selector)
	if setBody {
		if version != 0 {
			data := readConfig(storage)
			obj, err := api.Decode(data)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			jsonBase, err := api.FindJSONBase(obj)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			jsonBase.SetResourceVersion(version)
			data, err = api.Encode(obj)
			if err != nil {
				glog.Fatalf("error setting resource version: %v", err)
			}
			r.Body(data)
		} else {
			r.Body(readConfig(storage))
		}
	}
	result := r.Do()
	obj, err := result.Get()
	if err != nil {
		glog.Fatalf("Got request error: %v\n", err)
		return false
	}

	var printer kubecfg.ResourcePrinter
	switch {
	case *json:
		printer = &kubecfg.IdentityPrinter{}
	case *yaml:
		printer = &kubecfg.YAMLPrinter{}
	case len(*templateFile) > 0 || len(*templateStr) > 0:
		var data []byte
		if len(*templateFile) > 0 {
			var err error
			data, err = ioutil.ReadFile(*templateFile)
			if err != nil {
				glog.Fatalf("Error reading template %s, %v\n", *templateFile, err)
				return false
			}
		} else {
			data = []byte(*templateStr)
		}
		tmpl, err := template.New("output").Parse(string(data))
		if err != nil {
			glog.Fatalf("Error parsing template %s, %v\n", string(data), err)
			return false
		}
		printer = &kubecfg.TemplatePrinter{
			Template: tmpl,
		}
	default:
		printer = humanReadablePrinter()
	}

	if err = printer.PrintObj(obj, os.Stdout); err != nil {
		body, _ := result.Raw()
		glog.Fatalf("Failed to print: %v\nRaw received object:\n%#v\n\nBody received: %v", err, obj, string(body))
	}
	fmt.Print("\n")

	return true
}

func executeControllerRequest(method string, c *kube_client.Client) bool {
	parseController := func() string {
		if len(flag.Args()) != 2 {
			glog.Fatal("usage: kubecfg [OPTIONS] stop|rm|rollingupdate <controller>")
		}
		return flag.Arg(1)
	}

	var err error
	switch method {
	case "stop":
		err = kubecfg.StopController(parseController(), c)
	case "rm":
		err = kubecfg.DeleteController(parseController(), c)
	case "rollingupdate":
		err = kubecfg.Update(parseController(), c, *updatePeriod)
	case "run":
		if len(flag.Args()) != 4 {
			glog.Fatal("usage: kubecfg [OPTIONS] run <image> <replicas> <controller>")
		}
		image := flag.Arg(1)
		replicas, err2 := strconv.Atoi(flag.Arg(2))
		if err2 != nil {
			glog.Fatalf("Error parsing replicas: %v", err2)
		}
		name := flag.Arg(3)
		err = kubecfg.RunController(image, name, replicas, c, *portSpec, *servicePort)
	case "resize":
		args := flag.Args()
		if len(args) < 3 {
			glog.Fatal("usage: kubecfg resize <controller> <replicas>")
		}
		name := args[1]
		replicas, err2 := strconv.Atoi(args[2])
		if err2 != nil {
			glog.Fatalf("Error parsing replicas: %v", err2)
		}
		err = kubecfg.ResizeController(name, replicas, c)
	default:
		return false
	}
	if err != nil {
		glog.Fatalf("Error: %v", err)
	}
	return true
}

func humanReadablePrinter() *kubecfg.HumanReadablePrinter {
	printer := kubecfg.NewHumanReadablePrinter()
	// Add Handler calls here to support additional types
	return printer
}
<|endoftext|>


<commit_before>package profile

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"

	"github.com/pkg/errors"
)

type Profile struct {
	name       string
	cpuProfile *os.File
}

func New(name string) (*Profile, error) {
	// Open the CPU profile output.
	cpuProfile, err := os.Create(fmt.Sprintf("%s_cpu.prof", name))
	if err != nil {
		return nil, errors.Wrap(err, "unable to create CPU profile")
	}

	// Start CPU profiling.
	if err := pprof.StartCPUProfile(cpuProfile); err != nil {
		cpuProfile.Close()
		return nil, errors.Wrap(err, "unable to start CPU profile")
	}

	// Success.
	return &Profile{
		name:       name,
		cpuProfile: cpuProfile,
	}, nil
}

func (p *Profile) Finalize() error {
	// Close out the CPU profile.
	pprof.StopCPUProfile()
	if err := p.cpuProfile.Close(); err != nil {
		return errors.Wrap(err, "unable to close CPU profile")
	}

	// Run a GC cycle to update the heap profile statistics.
	runtime.GC()

	// Write a heap profile.
	heapProfile, err := os.Create(fmt.Sprintf("%s_heap.prof", p.name))
	if err != nil {
		return errors.Wrap(err, "unable to create heap profile")
	}
	if err := pprof.WriteHeapProfile(heapProfile); err != nil {
		heapProfile.Close()
		return errors.Wrap(err, "unable to write heap profile")
	}
	if err := heapProfile.Close(); err != nil {
		return errors.Wrap(err, "unable to close heap profile")
	}

	// Success.
	return nil
}
<commit_msg>Polished profile package.<commit_after>package profile

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"

	"github.com/pkg/errors"
)

// Profile manages a CPU and heap profile.
type Profile struct {
	// name is the name of the profile.
	name string
	// cpuProfile is the output file for the CPU profile.
	cpuProfile *os.File
}

// New creates a new profile instance. The profiling begins immediately.
func New(name string) (*Profile, error) {
	// Open the CPU profile output.
	cpuProfile, err := os.Create(fmt.Sprintf("%s_cpu.prof", name))
	if err != nil {
		return nil, errors.Wrap(err, "unable to create CPU profile")
	}

	// Start CPU profiling.
	if err := pprof.StartCPUProfile(cpuProfile); err != nil {
		cpuProfile.Close()
		return nil, errors.Wrap(err, "unable to start CPU profile")
	}

	// Success.
	return &Profile{
		name:       name,
		cpuProfile: cpuProfile,
	}, nil
}

// Finalize terminates a profile and writes its measurements to disk in the
// current working directory.
func (p *Profile) Finalize() error {
	// Close out the CPU profile.
	pprof.StopCPUProfile()
	if err := p.cpuProfile.Close(); err != nil {
		return errors.Wrap(err, "unable to close CPU profile")
	}

	// Run a GC cycle to update the heap profile statistics.
	runtime.GC()

	// Write a heap profile.
	heapProfile, err := os.Create(fmt.Sprintf("%s_heap.prof", p.name))
	if err != nil {
		return errors.Wrap(err, "unable to create heap profile")
	}
	if err := pprof.WriteHeapProfile(heapProfile); err != nil {
		heapProfile.Close()
		return errors.Wrap(err, "unable to write heap profile")
	}
	if err := heapProfile.Close(); err != nil {
		return errors.Wrap(err, "unable to close heap profile")
	}

	// Success.
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"knative.dev/serving/pkg/activator"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/queue"
)

const wantHost = "a-better-host.com"

func TestHandlerReqEvent(t *testing.T) {
	var httpHandler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get(activator.RevisionHeaderName) != "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		if r.Header.Get(activator.RevisionHeaderNamespace) != "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		if got, want := r.Host, wantHost; got != want {
			t.Errorf("Host header = %q, want: %q", got, want)
		}
		if got, want := r.Header.Get(network.OriginalHostHeader), ""; got != want {
			t.Errorf("%s header was preserved", network.OriginalHostHeader)
		}

		w.WriteHeader(http.StatusOK)
	}

	server := httptest.NewServer(httpHandler)
	serverURL, _ := url.Parse(server.URL)

	defer server.Close()
	proxy := httputil.NewSingleHostReverseProxy(serverURL)

	params := queue.BreakerParams{QueueDepth: 10, MaxConcurrency: 10, InitialCapacity: 10}
	breaker := queue.NewBreaker(params)
	reqChan := make(chan queue.ReqEvent, 10)
	h := handler(reqChan, breaker, proxy, func() bool { return true })

	writer := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)

	// Verify the Original host header processing.
	req.Host = "nimporte.pas"
	req.Header.Set(network.OriginalHostHeader, wantHost)

	req.Header.Set(network.ProxyHeaderName, activator.Name)
	h(writer, req)
	select {
	case e := <-reqChan:
		if e.EventType != queue.ProxiedIn {
			t.Errorf("Want: %v, got: %v\n", queue.ReqIn, e.EventType)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Timed out waiting for an event to be intercepted")
	}
}

func TestProbeHandler(t *testing.T) {
	testcases := []struct {
		name          string
		prober        func() bool
		wantCode      int
		wantBody      string
		requestHeader string
	}{{
		name:          "unexpected probe header",
		prober:        func() bool { return true },
		wantCode:      http.StatusBadRequest,
		wantBody:      fmt.Sprintf(badProbeTemplate, "test-probe"),
		requestHeader: "test-probe",
	}, {
		name:          "true probe function",
		prober:        func() bool { return true },
		wantCode:      http.StatusOK,
		wantBody:      queue.Name,
		requestHeader: queue.Name,
	}, {
		name:          "nil probe function",
		prober:        nil,
		wantCode:      http.StatusInternalServerError,
		wantBody:      "no probe",
		requestHeader: queue.Name,
	}, {
		name:          "false probe function",
		prober:        func() bool { return false },
		wantCode:      http.StatusServiceUnavailable,
		wantBody:      "container not ready",
		requestHeader: queue.Name,
	}}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			writer := httptest.NewRecorder()
			req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
			req.Header.Set(network.ProbeHeaderName, tc.requestHeader)

			h := handler(nil, nil, nil, tc.prober)
			h(writer, req)

			if got, want := writer.Code, tc.wantCode; got != want {
				t.Errorf("probe status = %v, want: %v", got, want)
			}
			if got, want := strings.TrimSpace(writer.Body.String()), tc.wantBody; got != want {
				// \r\n might be inserted, etc.
				t.Errorf("probe body = %q, want: %q, diff: %s", got, want, cmp.Diff(got, want))
			}
		})
	}
}

func TestCreateVarLogLink(t *testing.T) {
	dir, err := ioutil.TempDir("", "TestCreateVarLogLink")
	if err != nil {
		t.Errorf("Failed to created temporary directory: %v", err)
	}
	defer os.RemoveAll(dir)
	var env = config{
		ServingNamespace:   "default",
		ServingPod:         "service-7f97f9465b-5kkm5",
		UserContainerName:  "user-container",
		VarLogVolumeName:   "knative-var-log",
		InternalVolumePath: dir,
	}
	createVarLogLink(env)

	source := path.Join(dir, "default_service-7f97f9465b-5kkm5_user-container")
	want := "../knative-var-log"
	got, err := os.Readlink(source)
	if err != nil {
		t.Errorf("Failed to read symlink: %v", err)
	}
	if got != want {
		t.Errorf("Incorrect symlink = %q, want %q, diff: %s", got, want, cmp.Diff(got, want))
	}
}

func TestProbeQueueConnectionFailure(t *testing.T) {
	port := 12345 // some random port (that's not listening)

	if err := probeQueueHealthPath(port, 1); err == nil {
		t.Error("Expected error, got nil")
	}
}

func TestProbeQueueNotReady(t *testing.T) {
	queueProbed := false
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		queueProbed = true
		w.WriteHeader(http.StatusBadRequest)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	err = probeQueueHealthPath(port, 1)

	if diff := cmp.Diff(err.Error(), "probe returned not ready"); diff != "" {
		t.Errorf("Unexpected not ready message: %s", diff)
	}

	if !queueProbed {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueReady(t *testing.T) {
	queueProbed := false
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		queueProbed = true
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	if err = probeQueueHealthPath(port, 1); err != nil {
		t.Errorf("probeQueueHealthPath(%d, 1s) = %s", port, err)
	}

	if !queueProbed {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueTimeout(t *testing.T) {
	queueProbed := false
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		queueProbed = true
		time.Sleep(2 * time.Second)
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("failed to convert port(%s) to int", u.Port())
	}

	timeout := 1
	if err = probeQueueHealthPath(port, timeout); err == nil {
		t.Errorf("Expected probeQueueHealthPath(%d, %v) to return timeout error", port, timeout)
	}

	ts.Close()

	if !queueProbed {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueDelayedReady(t *testing.T) {
	count := 0
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if count < 9 {
			w.WriteHeader(http.StatusBadRequest)
			count++
			return
		}
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	timeout := 0
	if err := probeQueueHealthPath(port, timeout); err != nil {
		t.Errorf("probeQueueHealthPath(%d) = %s", port, err)
	}
}
<commit_msg>Fix race conditions in probe tests. (#4786)<commit_after>/*
Copyright 2019 The Knative Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"knative.dev/pkg/ptr"
	"knative.dev/serving/pkg/activator"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/queue"
)

const wantHost = "a-better-host.com"

func TestHandlerReqEvent(t *testing.T) {
	var httpHandler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get(activator.RevisionHeaderName) != "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		if r.Header.Get(activator.RevisionHeaderNamespace) != "" {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		if got, want := r.Host, wantHost; got != want {
			t.Errorf("Host header = %q, want: %q", got, want)
		}
		if got, want := r.Header.Get(network.OriginalHostHeader), ""; got != want {
			t.Errorf("%s header was preserved", network.OriginalHostHeader)
		}

		w.WriteHeader(http.StatusOK)
	}

	server := httptest.NewServer(httpHandler)
	serverURL, _ := url.Parse(server.URL)

	defer server.Close()
	proxy := httputil.NewSingleHostReverseProxy(serverURL)

	params := queue.BreakerParams{QueueDepth: 10, MaxConcurrency: 10, InitialCapacity: 10}
	breaker := queue.NewBreaker(params)
	reqChan := make(chan queue.ReqEvent, 10)
	h := handler(reqChan, breaker, proxy, func() bool { return true })

	writer := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)

	// Verify the Original host header processing.
	req.Host = "nimporte.pas"
	req.Header.Set(network.OriginalHostHeader, wantHost)

	req.Header.Set(network.ProxyHeaderName, activator.Name)
	h(writer, req)
	select {
	case e := <-reqChan:
		if e.EventType != queue.ProxiedIn {
			t.Errorf("Want: %v, got: %v\n", queue.ReqIn, e.EventType)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("Timed out waiting for an event to be intercepted")
	}
}

func TestProbeHandler(t *testing.T) {
	testcases := []struct {
		name          string
		prober        func() bool
		wantCode      int
		wantBody      string
		requestHeader string
	}{{
		name:          "unexpected probe header",
		prober:        func() bool { return true },
		wantCode:      http.StatusBadRequest,
		wantBody:      fmt.Sprintf(badProbeTemplate, "test-probe"),
		requestHeader: "test-probe",
	}, {
		name:          "true probe function",
		prober:        func() bool { return true },
		wantCode:      http.StatusOK,
		wantBody:      queue.Name,
		requestHeader: queue.Name,
	}, {
		name:          "nil probe function",
		prober:        nil,
		wantCode:      http.StatusInternalServerError,
		wantBody:      "no probe",
		requestHeader: queue.Name,
	}, {
		name:          "false probe function",
		prober:        func() bool { return false },
		wantCode:      http.StatusServiceUnavailable,
		wantBody:      "container not ready",
		requestHeader: queue.Name,
	}}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			writer := httptest.NewRecorder()
			req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
			req.Header.Set(network.ProbeHeaderName, tc.requestHeader)

			h := handler(nil, nil, nil, tc.prober)
			h(writer, req)

			if got, want := writer.Code, tc.wantCode; got != want {
				t.Errorf("probe status = %v, want: %v", got, want)
			}
			if got, want := strings.TrimSpace(writer.Body.String()), tc.wantBody; got != want {
				// \r\n might be inserted, etc.
				t.Errorf("probe body = %q, want: %q, diff: %s", got, want, cmp.Diff(got, want))
			}
		})
	}
}

func TestCreateVarLogLink(t *testing.T) {
	dir, err := ioutil.TempDir("", "TestCreateVarLogLink")
	if err != nil {
		t.Errorf("Failed to created temporary directory: %v", err)
	}
	defer os.RemoveAll(dir)
	var env = config{
		ServingNamespace:   "default",
		ServingPod:         "service-7f97f9465b-5kkm5",
		UserContainerName:  "user-container",
		VarLogVolumeName:   "knative-var-log",
		InternalVolumePath: dir,
	}
	createVarLogLink(env)

	source := path.Join(dir, "default_service-7f97f9465b-5kkm5_user-container")
	want := "../knative-var-log"
	got, err := os.Readlink(source)
	if err != nil {
		t.Errorf("Failed to read symlink: %v", err)
	}
	if got != want {
		t.Errorf("Incorrect symlink = %q, want %q, diff: %s", got, want, cmp.Diff(got, want))
	}
}

func TestProbeQueueConnectionFailure(t *testing.T) {
	port := 12345 // some random port (that's not listening)

	if err := probeQueueHealthPath(port, 1); err == nil {
		t.Error("Expected error, got nil")
	}
}

func TestProbeQueueNotReady(t *testing.T) {
	queueProbed := ptr.Int32(0)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(queueProbed, 1)
		w.WriteHeader(http.StatusBadRequest)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	err = probeQueueHealthPath(port, 1)

	if diff := cmp.Diff(err.Error(), "probe returned not ready"); diff != "" {
		t.Errorf("Unexpected not ready message: %s", diff)
	}

	if atomic.LoadInt32(queueProbed) == 0 {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueReady(t *testing.T) {
	queueProbed := ptr.Int32(0)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(queueProbed, 1)
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	if err = probeQueueHealthPath(port, 1); err != nil {
		t.Errorf("probeQueueHealthPath(%d, 1s) = %s", port, err)
	}

	if atomic.LoadInt32(queueProbed) == 0 {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueTimeout(t *testing.T) {
	queueProbed := ptr.Int32(0)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(queueProbed, 1)
		time.Sleep(2 * time.Second)
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("failed to convert port(%s) to int", u.Port())
	}

	timeout := 1
	if err = probeQueueHealthPath(port, timeout); err == nil {
		t.Errorf("Expected probeQueueHealthPath(%d, %v) to return timeout error", port, timeout)
	}

	ts.Close()

	if atomic.LoadInt32(queueProbed) == 0 {
		t.Errorf("Expected the queue proxy server to be probed")
	}
}

func TestProbeQueueDelayedReady(t *testing.T) {
	count := ptr.Int32(0)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if atomic.AddInt32(count, 1) < 9 {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		w.WriteHeader(http.StatusOK)
	}))

	defer ts.Close()

	u, err := url.Parse(ts.URL)
	if err != nil {
		t.Fatalf("%s is not a valid URL: %v", ts.URL, err)
	}

	port, err := strconv.Atoi(u.Port())
	if err != nil {
		t.Fatalf("Failed to convert port(%s) to int: %v", u.Port(), err)
	}

	timeout := 0
	if err := probeQueueHealthPath(port, timeout); err != nil {
		t.Errorf("probeQueueHealthPath(%d) = %s", port, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "flag"

  "../zeromq"
)

var (
  req_port    = flag.Int("req-port", 9797, "what Socket PORT to run at")
  rep_port    = flag.Int("rep-port", 9898, "what Socket PORT to run at")
)

func main(){
  flag.Parse()
  fmt.Printf("client ZeroMQ REP/REQ... at %d, %d", req_port, rep_port)

  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "myname", "anon")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "myname", "anonymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "delete", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "myname")
}
<commit_msg>updated zxtra/go0mq_client.go to use NS and TSDS dbtasks<commit_after>package main

import (
  "fmt"
  "flag"

  "../zeromq"
)

var (
  req_port    = flag.Int("req-port", 9797, "what Socket PORT to run at")
  rep_port    = flag.Int("rep-port", 9898, "what Socket PORT to run at")
)

func main(){
  flag.Parse()
  fmt.Printf("client ZeroMQ REP/REQ... at %d, %d", req_port, rep_port)

  fmt.Println("Checking out levigo based storage...")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "default", "myname", "anon")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "default", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "default", "myname", "anonymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "default", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "delete", "default", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "default", "myname")

  fmt.Println("Checking out levigoNS based storage...")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "ns", "myname:last:first", "anon")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "ns", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "ns", "myname:last", "ymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "ns", "myname", "anonymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "ns", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "delete", "ns", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "ns", "myname")

  fmt.Println("Checking out levigoTSDS based storage...")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "tsds-now", "myname:last:first", "anon")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "tsds", "myname:last:first", "2014", "2", "10", "9", "8", "7", "anon")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "tsds", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "tsds-now", "myname:last", "ymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "tsds", "myname", "2014", "2", "10", "9", "18", "37", "anonymous")
  abkzeromq.ZmqReq(*req_port, *rep_port, "push", "tsds", "myname", "2014", "2", "10", "5", "28", "57", "untitles")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "tsds", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "tsds", "myname:2014:February:10")
  abkzeromq.ZmqReq(*req_port, *rep_port, "delete", "tsds", "myname")
  abkzeromq.ZmqReq(*req_port, *rep_port, "read", "tsds", "myname")
}
<|endoftext|>


<commit_before>package numeric

import (
	"github.com/ready-steady/linear/matrix"
)

// Compute calculates the temperature profile corresponding to a power profile.
// The power profile is specified by a function func(time float64, power
// []float64) evaluating the power dissipation at an arbitrary time moment. The
// time moments for which the temperature profile is computed are specified by
// the time array; see the corresponding ODE solver for further details.
//
// http://godoc.org/github.com/ready-steady/numeric/integration/ode#DormandPrince.Compute
func (t *Temperature) Compute(power func(float64, []float64),
	time []float64) ([]float64, []float64, error) {

	nc, nn := t.nc, t.nn

	A, B := t.system.A, t.system.B

	P := make([]float64, nc)

	dSdt := func(time float64, S, dSdt []float64) {
		matrix.Multiply(A, S, dSdt, nn, nn, 1)
		power(time, P)
		for i := uint(0); i < nc; i++ {
			dSdt[i] += B[i] * P[i]
		}
	}

	S, time, err := t.integrator.Compute(dSdt, make([]float64, nn), time)
	if err != nil {
		return nil, nil, err
	}

	ns := uint(len(time))

	Q := make([]float64, ns*nc)
	for i := uint(0); i < nc; i++ {
		for j := uint(0); j < ns; j++ {
			Q[j*nc+i] = S[j*nn+i] + t.system.Qamb
		}
	}

	return Q, time, nil
}
<commit_msg>temperature/numeric: cosmetic changes<commit_after>package numeric

import (
	"github.com/ready-steady/linear/matrix"
)

// Compute calculates the temperature profile corresponding to a power profile.
// The power profile is specified by a function func(time float64, power
// []float64) evaluating the power dissipation at an arbitrary time moment. The
// time moments for which the temperature profile is computed are specified by
// the time array; see the corresponding ODE solver for further details.
//
// http://godoc.org/github.com/ready-steady/numeric/integration/ode#DormandPrince.Compute
func (t *Temperature) Compute(power func(float64, []float64),
	time []float64) ([]float64, []float64, error) {

	nc, nn := t.nc, t.nn

	A, B := t.system.A, t.system.B
	P := make([]float64, nc)

	dSdt := func(t float64, S, dSdt []float64) {
		matrix.Multiply(A, S, dSdt, nn, nn, 1)
		power(t, P)
		for i := uint(0); i < nc; i++ {
			dSdt[i] += B[i] * P[i]
		}
	}

	S, time, err := t.integrator.Compute(dSdt, make([]float64, nn), time)
	if err != nil {
		return nil, nil, err
	}

	ns := uint(len(time))

	Q, Qamb := make([]float64, ns*nc), t.system.Qamb
	for i := uint(0); i < nc; i++ {
		for j := uint(0); j < ns; j++ {
			Q[j*nc+i] = S[j*nn+i] + Qamb
		}
	}

	return Q, time, nil
}
<|endoftext|>