From 33d6b1cd1277e68afe73e36770f555ff38b1d6d7 Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmyzelinskie@gmail.com>
Date: Wed, 2 Mar 2016 20:59:01 -0500
Subject: [PATCH] implement prometheus server

---
 cmd/chihaya/main.go             |  1 +
 config_example.yaml             | 13 ++++-
 server/prometheus/prometheus.go | 93 +++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+), 2 deletions(-)
 create mode 100644 server/prometheus/prometheus.go

diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go
index 60f5061..3d59721 100644
--- a/cmd/chihaya/main.go
+++ b/cmd/chihaya/main.go
@@ -16,6 +16,7 @@ import (
 	"github.com/chihaya/chihaya/tracker"
 
 	_ "github.com/chihaya/chihaya/server/http"
+	_ "github.com/chihaya/chihaya/server/prometheus"
 	_ "github.com/chihaya/chihaya/server/store"
 	_ "github.com/chihaya/chihaya/server/store/memory"
 	_ "github.com/chihaya/chihaya/server/store/middleware/client"
diff --git a/config_example.yaml b/config_example.yaml
index 45874a2..664f2e6 100644
--- a/config_example.yaml
+++ b/config_example.yaml
@@ -7,10 +7,12 @@ chihaya:
     announce: 10m
     min_announce: 5m
     announce_middleware:
+      # These are currently fake values
       - prometheus
       - store_client_validation
       - store_create_on_announce
     scrape_middleware:
+      # These are currently fake values
       - prometheus
       - store_client_validation
 
@@ -28,13 +30,20 @@ chihaya:
           gcAfter: 30m
           shards: 1
 
-    - name: http
+    - name: prometheus
       config:
         addr: localhost:6881
+        shutdown_timeout: 10s
+        read_timeout: 10s
+        write_timeout: 10s
+
+    - name: http
+      config:
+        addr: localhost:6882
         request_timeout: 10s
         read_timeout: 10s
         write_timeout: 10s
 
     - name: udp
       config:
-        addr: localhost:6882
+        addr: localhost:6883
diff --git a/server/prometheus/prometheus.go b/server/prometheus/prometheus.go
new file mode 100644
index 0000000..f92c3c9
--- /dev/null
+++ b/server/prometheus/prometheus.go
@@ -0,0 +1,93 @@
+// Copyright 2016 The Chihaya Authors. All rights reserved.
+// Use of this source code is governed by the BSD 2-Clause license,
+// which can be found in the LICENSE file.
+
+// Package prometheus implements a chihaya Server for serving metrics to
+// Prometheus.
+package prometheus
+
+import (
+	"errors"
+	"net/http"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/tylerb/graceful"
+	"gopkg.in/yaml.v2"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/server"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+func init() {
+	server.Register("prometheus", constructor)
+}
+
+func constructor(srvcfg *chihaya.ServerConfig, tkr *tracker.Tracker) (server.Server, error) {
+	cfg, err := NewServerConfig(srvcfg)
+	if err != nil {
+		return nil, errors.New("prometheus: invalid config: " + err.Error())
+	}
+
+	return &Server{
+		cfg: cfg,
+	}, nil
+}
+
+// ServerConfig represents the configuration options for a
+// PrometheusServer.
+type ServerConfig struct {
+	Addr            string        `yaml:"addr"`
+	ShutdownTimeout time.Duration `yaml:"shutdown_timeout"`
+	ReadTimeout     time.Duration `yaml:"read_timeout"`
+	WriteTimeout    time.Duration `yaml:"write_timeout"`
+}
+
+// NewServerConfig marshals a chihaya.ServerConfig and unmarshals it
+// into a more specific prometheus ServerConfig.
+func NewServerConfig(srvcfg *chihaya.ServerConfig) (*ServerConfig, error) {
+	bytes, err := yaml.Marshal(srvcfg.Config)
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg ServerConfig
+	err = yaml.Unmarshal(bytes, &cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
+// Server implements a chihaya Server for serving metrics to Prometheus.
+type Server struct {
+	cfg     *ServerConfig
+	grace   *graceful.Server
+	stopped bool
+}
+
+var _ server.Server = &Server{}
+
+func (s *Server) Start() {
+	s.grace = &graceful.Server{
+		Server: &http.Server{
+			Addr:         s.cfg.Addr,
+			Handler:      prometheus.Handler(),
+			ReadTimeout:  s.cfg.ReadTimeout,
+			WriteTimeout: s.cfg.WriteTimeout,
+		},
+		Timeout:          s.cfg.ShutdownTimeout,
+		NoSignalHandling: true,
+	}
+}
+
+func (s *Server) Stop() {
+	s.grace.Stop(s.cfg.ShutdownTimeout)
+	stopChan := s.grace.StopChan()
+
+	// Block until the graceful server shuts down and closes this channel.
+	for range stopChan {
+	}
+}