Added the -cache-limit option to limit cache sizes.

cmd/cesium-terrain-server/limitopt.go

@@ -0,0 +1,96 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"github.com/geo-data/cesium-terrain-server/handlers"
+	"strconv"
+)
+
+// Adapted from <https://golang.org/doc/effective_go.html#constants>.
+type ByteSize float64
+
+const (
+	_           = iota // ignore first value by assigning to blank identifier
+	KB ByteSize = 1 << (10 * iota)
+	MB
+	GB
+	TB
+)
+
+func (b ByteSize) String() string {
+	switch {
+	case b >= TB:
+		return fmt.Sprintf("%.2fTB", b/TB)
+	case b >= GB:
+		return fmt.Sprintf("%.2fGB", b/GB)
+	case b >= MB:
+		return fmt.Sprintf("%.2fMB", b/MB)
+	case b >= KB:
+		return fmt.Sprintf("%.2fkB", b/KB)
+	}
+	return fmt.Sprintf("%.2fB", b)
+}
+
+func ParseByteSize(size string) (bytes ByteSize, err error) {
+	defer func() {
+		if bytes < 0 {
+			err = errors.New("size cannot be negative")
+		}
+	}()
+
+	val, err := strconv.ParseFloat(size, 64)
+	if err == nil {
+		bytes = ByteSize(val)
+		return
+	}
+
+	if len(size) < 3 {
+		err = errors.New("the size must be specified as a suffix e.g 5MB")
+		return
+	}
+
+	val, err = strconv.ParseFloat(size[:len(size)-2], 64)
+	if err != nil {
+		return
+	}
+	bytes = ByteSize(val)
+
+	suffix := size[len(size)-2:]
+	switch suffix {
+	case "TB":
+		bytes *= TB
+	case "GB":
+		bytes *= GB
+	case "MB":
+		bytes *= MB
+	case "KB":
+		bytes *= KB
+	default:
+		err = errors.New("bad size suffix: " + suffix)
+	}
+	return
+}
+
+type LimitOpt struct {
+	Value handlers.Bytes
+}
+
+func NewLimitOpt() *LimitOpt {
+	return &LimitOpt{}
+}
+
+func (this *LimitOpt) String() string {
+	return ByteSize(this.Value).String()
+}
+
+func (this *LimitOpt) Set(size string) error {
+	byteSize, err := ParseByteSize(size)
+	if err != nil {
+		return err
+	}
+
+	this.Value = handlers.Bytes(byteSize)
+
+	return nil
+}

cmd/cesium-terrain-server/logopt.go

@@ -0,0 +1,45 @@
+package main
+
+import (
+	"errors"
+	"github.com/geo-data/cesium-terrain-server/log"
+)
+
+type LogOpt struct {
+	Priority log.Priority
+}
+
+func NewLogOpt() *LogOpt {
+	return &LogOpt{
+		Priority: log.LOG_NOTICE,
+	}
+}
+
+func (this *LogOpt) String() string {
+	switch this.Priority {
+	case log.LOG_CRIT:
+		return "crit"
+	case log.LOG_ERR:
+		return "err"
+	case log.LOG_NOTICE:
+		return "notice"
+	default:
+		return "debug"
+	}
+}
+
+func (this *LogOpt) Set(level string) error {
+	switch level {
+	case "crit":
+		this.Priority = log.LOG_CRIT
+	case "err":
+		this.Priority = log.LOG_ERR
+	case "notice":
+		this.Priority = log.LOG_NOTICE
+	case "debug":
+		this.Priority = log.LOG_DEBUG
+	default:
+		return errors.New("choose one of crit, err, notice, debug")
+	}
+	return nil
+}

cmd/cesium-terrain-server/main.go

@@ -2,7 +2,6 @@
 package main
 
 import (
-	"errors"
 	"flag"
 	"fmt"
 	myhandlers "github.com/geo-data/cesium-terrain-server/handlers"
@@ -15,45 +14,6 @@ import (
 	"os"
 )
 
-type LogOpt struct {
-	Priority log.Priority
-}
-
-func NewLogOpt() *LogOpt {
-	return &LogOpt{
-		Priority: log.LOG_NOTICE,
-	}
-}
-
-func (this *LogOpt) String() string {
-	switch this.Priority {
-	case log.LOG_CRIT:
-		return "crit"
-	case log.LOG_ERR:
-		return "err"
-	case log.LOG_NOTICE:
-		return "notice"
-	default:
-		return "debug"
-	}
-}
-
-func (this *LogOpt) Set(level string) error {
-	switch level {
-	case "crit":
-		this.Priority = log.LOG_CRIT
-	case "err":
-		this.Priority = log.LOG_ERR
-	case "notice":
-		this.Priority = log.LOG_NOTICE
-	case "debug":
-		this.Priority = log.LOG_DEBUG
-	default:
-		return errors.New("choose one of crit, err, notice, debug")
-	}
-	return nil
-}
-
 func main() {
 	port := flag.Uint("port", 8000, "the port on which the server listens")
 	tilesetRoot := flag.String("dir", ".", "the root directory under which tileset directories reside")
@@ -63,6 +23,9 @@ func main() {
 	noRequestLog := flag.Bool("no-request-log", false, "do not log client requests for resources")
 	logging := NewLogOpt()
 	flag.Var(logging, "log-level", "level at which logging occurs. One of crit, err, notice, debug")
+	limit := NewLimitOpt()
+	limit.Set("1MB")
+	flag.Var(limit, "cache-limit", `the memory size in bytes beyond which resources are not cached. Other memory units can be specified by suffixing the number with kB, MB, GB or TB`)
 	flag.Parse()
 
 	// Set the logging
@@ -82,7 +45,7 @@ func main() {
 	handler := myhandlers.AddCorsHeader(r)
 	if len(*memcached) > 0 {
 		log.Debug(fmt.Sprintf("memcached enabled for all resources: %s", *memcached))
-		handler = myhandlers.NewCache(*memcached, handler)
+		handler = myhandlers.NewCache(*memcached, handler, limit.Value, myhandlers.NewLimit)
 	}
 
 	if *noRequestLog == false {

handlers/cache.go

@@ -11,12 +11,16 @@ import (
 type Cache struct {
 	mc      *memcache.Client
 	handler http.Handler
+	Limit   Bytes
+	limiter LimiterFactory
 }
 
-func NewCache(connstr string, handler http.Handler) http.Handler {
+func NewCache(connstr string, handler http.Handler, limit Bytes, limiter LimiterFactory) http.Handler {
 	return &Cache{
 		mc:      memcache.New(connstr),
 		handler: handler,
+		Limit:   limit,
+		limiter: limiter,
 	}
 }
 
@@ -33,12 +37,30 @@ func (this *Cache) generateKey(r *http.Request) string {
 }
 
 func (this *Cache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	var limiter ResponseLimiter
+	var recorder http.ResponseWriter
 	rec := NewRecorder()
 
-	// Write to both the recorder and original writer
-	tee := MultiWriter(w, rec)
+	// If a limiter is provided, wrap the recorder with it.
+	if this.limiter != nil {
+		limiter = this.limiter(rec, this.Limit)
+		recorder = limiter
+	} else {
+		recorder = rec
+	}
+
+	// Write to both the recorder and original writer.
+	tee := MultiWriter(w, recorder)
 	this.handler.ServeHTTP(tee, r)
 
+	// If the cache limit has been exceeded, don't proceed to cache the
+	// response.
+	if limiter != nil && limiter.LimitExceeded() {
+		log.Debug(fmt.Sprintf("cache limit exceeded for %s", r.URL.String()))
+		return
+	}
+
+	// Cache the response.
 	key := this.generateKey(r)
 	log.Debug(fmt.Sprintf("setting key: %s", key))
 	if err := this.mc.Set(&memcache.Item{Key: key, Value: rec.Body.Bytes()}); err != nil {

handlers/handlers.go

@@ -2,6 +2,15 @@ package handlers
 
 import "net/http"
 
+type Bytes uint64
+
+type ResponseLimiter interface {
+	http.ResponseWriter
+	LimitExceeded() bool
+}
+
+type LimiterFactory func(writer http.ResponseWriter, limit Bytes) ResponseLimiter
+
 // Return HTTP middleware which allows CORS requests from any domain
 func AddCorsHeader(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

handlers/limiter.go

@@ -0,0 +1,54 @@
+package handlers
+
+import (
+	"net/http"
+)
+
+// ResponseLimiter is an implementation of ResponseLimiter that wraps a
+// http.ResponseWriter and limits how much is written to it. If an attempt is
+// made to write more bytes than the limit it silently fails.
+type LimitResponse struct {
+	Writer   http.ResponseWriter
+	Limit    Bytes // the maximum number of bytes that can be written to Writer
+	written  Bytes // the number of bytes already sent to Writer
+	exceeded bool  // has the limit been exceeded?
+}
+
+// NewRecorder returns an initialized LimitResponse. This implements the
+// LimiterFactory function type.
+func NewLimit(writer http.ResponseWriter, limit Bytes) ResponseLimiter {
+	return &LimitResponse{
+		Writer: writer,
+		Limit:  limit,
+	}
+}
+
+func (this *LimitResponse) LimitExceeded() bool {
+	return this.exceeded
+}
+
+// Header returns the response headers.
+func (this *LimitResponse) Header() http.Header {
+	return this.Writer.Header()
+}
+
+// Write always succeeds and writes to this.Body, if not nil.
+func (this *LimitResponse) Write(buf []byte) (bytes int, err error) {
+	if this.exceeded {
+		return
+	}
+
+	if (Bytes(len(buf)) + this.written) > this.Limit {
+		this.exceeded = true
+		return
+	}
+
+	bytes, err = this.Writer.Write(buf)
+	this.written += Bytes(bytes)
+	return
+}
+
+// WriteHeader sets this.Code.
+func (this *LimitResponse) WriteHeader(code int) {
+	this.Writer.WriteHeader(code)
+}