Factor common code into Graphite package

cmd/influxd/run.go

@@ -159,25 +159,18 @@ func Run(config *Config, join, version string, logWriter *os.File) (*messaging.B
 			parser.Separator = c.NameSeparatorString()
 			parser.LastEnabled = c.LastEnabled()
 
-			// Start the relevant server.
-			if strings.ToLower(c.Protocol) == "tcp" {
-				g := graphite.NewTCPServer(parser, s)
-				g.Database = c.Database
-				g.SetLogOutput(logWriter)
-				err := g.ListenAndServe(c.ConnectionString(config.BindAddress))
-				if err != nil {
-					log.Printf("failed to start TCP Graphite Server: %v\n", err.Error())
-				}
-			} else if strings.ToLower(c.Protocol) == "udp" {
-				g := graphite.NewUDPServer(parser, s)
-				g.Database = c.Database
-				g.SetLogOutput(logWriter)
-				err := g.ListenAndServe(c.ConnectionString(config.BindAddress))
-				if err != nil {
-					log.Printf("failed to start UDP Graphite Server: %v\n", err.Error())
-				}
-			} else {
-				log.Fatalf("unrecognized Graphite Server protocol %s", c.Protocol)
+			// Spin up the server.
+			var g graphite.Server
+			g, err := graphite.NewServer(c.Protocol, parser, s)
+			if err != nil {
+				log.Fatalf("failed to initialize %s Graphite server: %s", c.Protocol, err.Error())
+			}
+
+			g.SetDatabase(c.Database)
+			g.SetLogOutput(logWriter)
+			err = g.ListenAndServe(c.ConnectionString(config.BindAddress))
+			if err != nil {
+				log.Fatalf("failed to start %s Graphite server: %s", c.Protocol, err.Error())
 			}
 		}
 	}

graphite/graphite.go

@@ -3,6 +3,7 @@ package graphite
 import (
 	"errors"
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 	"time"
@@ -41,6 +42,26 @@ type SeriesWriter interface {
 	DatabaseExists(string) bool
 }
 
+// Server defines the interface all Graphite servers support.
+type Server interface {
+	SetLogOutput(w io.Writer)
+	SetDatabase(string)
+	ListenAndServe(iface string) error
+	Protocol() string
+}
+
+// NewServer return a Graphite server for the given protocol, using the given parser
+// and series writer.
+func NewServer(protocol string, p *Parser, s SeriesWriter) (Server, error) {
+	if strings.ToLower(protocol) == "tcp" {
+		return NewTCPServer(p, s), nil
+	} else if strings.ToLower(protocol) == "udp" {
+		return NewUDPServer(p, s), nil
+	} else {
+		return nil, fmt.Errorf("unrecognized Graphite Server protocol %s", protocol)
+	}
+}
+
 // Parser encapulates a Graphite Parser.
 type Parser struct {
 	Separator   string

graphite/graphite_tcp.go

@@ -32,6 +32,16 @@ func (s *TCPServer) SetLogOutput(w io.Writer) {
 	s.Logger = log.New(w, "[graphite] ", log.LstdFlags)
 }
 
+// SetDatabase sets database for all Graphite log output.
+func (s *TCPServer) SetDatabase(database string) {
+	s.Database = database
+}
+
+// Protocol returns a string version of the supported protocol.
+func (s *TCPServer) Protocol() string {
+	return "tcp"
+}
+
 // ListenAndServe instructs the TCPServer to start processing Graphite data
 // on the given interface. iface must be in the form host:port
 func (t *TCPServer) ListenAndServe(iface string) error {

graphite/graphite_udp.go

@@ -36,6 +36,16 @@ func (s *UDPServer) SetLogOutput(w io.Writer) {
 	s.Logger = log.New(w, "[graphite] ", log.LstdFlags)
 }
 
+// SetDatabase sets database for all Graphite log output.
+func (s *UDPServer) SetDatabase(database string) {
+	s.Database = database
+}
+
+// Protocol returns a string version of the supported protocol.
+func (s *UDPServer) Protocol() string {
+	return "udp"
+}
+
 // ListenAndServer instructs the UDPServer to start processing Graphite data
 // on the given interface. iface must be in the form host:port.
 func (u *UDPServer) ListenAndServe(iface string) error {