metrorama/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"
)

var (
	listen     = flag.String("l", "0.0.0.0:8080", "Listen address")
	rate       = flag.Int("r", 500, "Sample rate (ms)")
	timeout    = flag.Int("t", 10, "Total timeout (s)")
	minRate    = 500
	maxTimeout = 600
)

// TODO: bench test with pointer
func getMetrics() (Metric, error) {
	log.Println("Started getMetrics")
	m := Metric{}
	err := m.getMemory()
	if err != nil {
		log.Printf("Cannot get memory: %s\n", err)
		return m, err
	}

	err = m.getLoadAvg()
	if err != nil {
		log.Printf("Cannot get load average: %s\n", err)
		return m, err
	}

	err = m.getCpuMetric()
	if err != nil {
		log.Printf("Cannot get CPU metrics: %s\n", err)
		return m, err
	}

	err = m.getDiskMetrics()
	if err != nil {
		log.Printf("Cannot get disk metrics: %s\n", err)
		return m, err
	}

	err = m.getNetworkMetrics()
	if err != nil {
		log.Printf("Cannot get network metrics: %s\n", err)
		return m, err
	}

	return m, nil
}

func (t *Test) collectData() {
	timeo := time.After(t.Timeout)
	t.Status = "STARTED"
	for {
		select {
		case tick := <-t.ticker.C:
			t.Status = "ONGOING"

			m, err := getMetrics()
			if err != nil {
				log.Printf("Cannot collect metrics: %s\n", err)
				return
			}
			t.Result = append(t.Result, m)
			log.Printf("%v: %v\n", tick, m)
			//t.Result += fmt.Sprintf("%v: Collecting data\n", tick)
		case <-t.Done:
			log.Printf("DONE\n")
			t.Status = "COMPLETED"
			t.Unlock()
			return
		case <-timeo:
			log.Printf("TIMEOUT\n")
			t.Status = "TIMEOUT"
			t.Unlock()
			return
		}
	}
	return // really needed?
}

func (t *Test) Start(w http.ResponseWriter, req *http.Request) {
	// Check if another collectData is running (better: check sync.Mutex)
	if t.Status == "ONGOING" {
		http.Error(w, "Another test is currently ongoing, stop or try again later", 409) //better error code
		return
	}
	// Check parameters
	qRate := req.URL.Query().Get("rate")
	if qRate != "" {
		qRate, err := strconv.Atoi(qRate)
		if err != nil {
			log.Printf("Error while converting rate: %s\n", err.Error())
			http.Error(w, "Error: error while converting rate", 400)
			return
		}
		if qRate < minRate {
			log.Printf("Requested rate is too low: %d (min: %d)\n", qRate, minRate)
			http.Error(w, "Error: requested rate too low", 400)
			return
		}
		t.Rate = qRate
	}

	qTimeout := req.URL.Query().Get("timeout")
	if qTimeout != "" {
		qTimeout, err := strconv.Atoi(qTimeout)
		if err != nil {
			log.Printf("Error while converting timeout: %s\n", err.Error())
			http.Error(w, "Error: error while converting timeout", 400)
			return
		}
		if qTimeout > maxTimeout {
			log.Printf("Requested timeout is too high: %d (max: %d)\n", qTimeout, maxTimeout)
			http.Error(w, "Error: requested timeout too high", 400)
			return
		}
		t.Timeout = time.Duration(qTimeout) * time.Second
	}

	qName := req.URL.Query().Get("name")
	if qName == "" {
		t.Name = fmt.Sprintf("test-%d", time.Now().Local().Unix())
	} else {
		t.Name = qName
	}
	// Check if previous test is present
	if len(t.Result) == 0 {
		log.Println("Cleaning previous test data")
		t.Result = []Metric{}
	}
	// Lock to avoid multiple start
	t.Lock()
	fmt.Fprintf(w, "Test name: %s, rate: %d (ms), timeout: %d (s)\n", t.Name, t.Rate, t.Timeout/1000000000)
	endTime := time.Now().Local().Add(t.Timeout)
	fmt.Fprintf(w, "Test will end @%v\n", endTime.Format(time.RFC1123Z))
	go t.collectData()
}

func (t *Test) Stop(w http.ResponseWriter, req *http.Request) {
	if t.Status != "ONGOING" {
		e := fmt.Sprintf("Cannot stop test in %s status: try start first", t.Status)
		http.Error(w, e, 400)
		return
	}

	t.ticker.Stop()
	t.Done <- true
	fmt.Fprintf(w, "Done\n")
	fmt.Fprintf(w, "%v\n", t.Result)
}

func (t *Test) getStatus(w http.ResponseWriter, req *http.Request) {
	qStream := req.URL.Query().Get("stream")
	qFormat := req.URL.Query().Get("format")

	flusher, ok := w.(http.Flusher)
	if !ok {
		if qStream == "true" {
			log.Printf("Http server does not support flusher\n")
			http.Error(w, "Server does not support flusher",
				http.StatusInternalServerError)
			return
		}
	}

	if qStream == "true" {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
	}

	if qFormat == "json" {
		for {
			b, err := json.Marshal(t.Result)
			if err != nil {
				log.Printf("Cannot marshal json: %s\n", err)
				http.Error(w, err.Error(), 500)
				return
			}

			fmt.Fprintf(w, "%s\n", b)

			if qStream == "true" {
				time.Sleep(time.Duration(t.Rate) * time.Millisecond)
				flusher.Flush()
			} else {
				return
			}
		}
	} else {
		for {
			fmt.Fprintf(w, "%+ v\n", t.Result)
			if qStream == "true" {
				time.Sleep(time.Duration(t.Rate) * time.Millisecond)
				flusher.Flush()
			} else {
				return
			}
		}
		fmt.Fprintf(w, "Test %s, status: %s\n", t.Name, t.Status)
	}
}

func newTest() *Test {
	log.Println("Creating new test")
	// Not checking if rate and timeout are in value range because are set on CLI
	t := Test{
		ticker:  time.NewTicker(time.Duration(*rate) * time.Millisecond),
		Done:    make(chan bool),
		Rate:    *rate,
		Timeout: time.Duration(*timeout) * time.Second,
		Status:  "CREATED",
	}
	return &t
}

func main() {

	flag.Parse()

	t := newTest()
	log.Printf("%+v\n", t)

	http.HandleFunc("/start", t.Start)
	http.HandleFunc("/status", t.getStatus)
	http.HandleFunc("/stop", t.Stop)

	log.Printf("Listening on %s\n", *listen)
	log.Fatal(http.ListenAndServe(*listen, nil))
}