Added basic statistics under /stats

OrbitalBeat.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"encoding/base64"
 	"encoding/hex"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"git.gammaspectra.live/S.O.N.G/MakyuuIchaival"
@@ -20,8 +21,10 @@ import (
 	"net/url"
 	"os"
 	"path"
+	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -33,6 +36,7 @@ type ContentCacheEntry struct {
 var dbHandle *sql.DB
 var sha256Statement *sql.Stmt
 var md5Statement *sql.Stmt
+var statsStatement *sql.Stmt
 var fdlimit int
 
 var objectCacheMutex sync.RWMutex
@@ -42,6 +46,41 @@ var trustedPublicKeys []ed25519.PublicKey
 
 var debugOutput = false
 
+type statistics struct {
+	Served         servedStatistics         `json:"served"`
+	ContentCache   contentCacheStatistics   `json:"content_cache"`
+	SignatureCache signatureCacheStatistics `json:"signature_cache"`
+	HttpCode       httpCodeStatistics       `json:"http_code"`
+}
+
+type servedStatistics struct {
+	TotalBytesServed           uint64 `json:"bytes"`
+	TotalHashesServed          uint64 `json:"hashes"`
+	TotalContentRequestsServed uint64 `json:"content_requests"`
+}
+
+type contentCacheStatistics struct {
+	ContentCacheHit  uint64 `json:"cache_hit"`
+	ContentCacheMiss uint64 `json:"cache_miss"`
+}
+
+type signatureCacheStatistics struct {
+	SignatureCacheHit  uint64 `json:"cache_hit"`
+	SignatureCacheMiss uint64 `json:"cache_miss"`
+}
+
+type httpCodeStatistics struct {
+	Code200 uint64 `json:"200"`
+	Code204 uint64 `json:"204"`
+	Code302 uint64 `json:"302"`
+	Code400 uint64 `json:"400"`
+	Code403 uint64 `json:"403"`
+	Code404 uint64 `json:"404"`
+	Code501 uint64 `json:"501"`
+}
+
+var globalStatistics statistics
+
 func getFirstValidContentEntry(entries *[]ContentEntry) *ContentEntry {
 	for _, entry := range *entries {
 		stat, err := os.Stat(entry.Path)
@@ -101,19 +140,13 @@ func GetMimeTypeFromExtension(ext string) string {
 		//Image types
 		case "png":
 			return "image/png"
-		case "jfif":
-			fallthrough
-		case "jpeg":
-			fallthrough
-		case "jpg":
+		case "jpg", "jpeg", "jfif":
 			return "image/jpeg"
 		case "gif":
 			return "image/gif"
 		case "svg":
 			return "image/svg+xml"
-		case "tiff":
-			fallthrough
-		case "tif":
+		case "tif", "tiff":
 			return "image/tiff"
 		case "webp":
 			return "image/webp"
@@ -178,8 +211,10 @@ func handleQueryRequest(ctx httputils.RequestContext, identifier cid.Cid, extraA
 		if entry != nil {
 			cacheEntry = getCacheEntryForContentEntry(entry, identifier)
 		}
+		atomic.AddUint64(&globalStatistics.ContentCache.ContentCacheMiss, 1)
 		ctx.AddTimingInformational("ec", "Content Cache MISS")
 	} else {
+		atomic.AddUint64(&globalStatistics.ContentCache.ContentCacheHit, 1)
 		ctx.AddTimingInformational("ec", "Content Cache HIT")
 	}
 	pTime := cTime
@@ -207,10 +242,12 @@ func handleQueryRequest(ctx httputils.RequestContext, identifier cid.Cid, extraA
 				kind = "md5"
 			}
 			if kind != "" {
+				atomic.AddUint64(&globalStatistics.HttpCode.Code302, 1)
 				ctx.DoRedirect(fmt.Sprintf("%s/%s/%s/%s", origin, kind, hex.EncodeToString(mh.Digest), strings.Join(extraArguments, "/")), http.StatusFound)
 				return
 			}
 		}
+		atomic.AddUint64(&globalStatistics.HttpCode.Code404, 1)
 		ctx.SetResponseCode(http.StatusNotFound)
 		return
 	}
@@ -219,6 +256,7 @@ func handleQueryRequest(ctx httputils.RequestContext, identifier cid.Cid, extraA
 
 	file, err := os.Open(cacheEntry.Entry.Path)
 	if err != nil {
+		atomic.AddUint64(&globalStatistics.HttpCode.Code404, 1)
 		ctx.SetResponseCode(http.StatusNotFound)
 		return
 	}
@@ -244,7 +282,41 @@ func handleQueryRequest(ctx httputils.RequestContext, identifier cid.Cid, extraA
 		ctx.SetResponseHeader("Content-Type", mime)
 	}
 
-	ctx.ServeStream(httputils.NewStreamFromFile(file))
+	atomic.AddUint64(&globalStatistics.Served.TotalContentRequestsServed, 1)
+	ss := &statsStream{
+		stream: httputils.NewStreamFromFile(file),
+	}
+	ctx.ServeStream(ss)
+	runtime.SetFinalizer(ss, (*statsStream).Close)
+}
+
+type statsStream struct {
+	stream    httputils.DefinedStream
+	byteCount int64
+}
+
+func (ss *statsStream) Read(p []byte) (n int, err error) {
+	ss.byteCount += int64(len(p))
+	return ss.stream.Read(p)
+}
+
+func (ss *statsStream) Seek(offset int64, whence int) (int64, error) {
+	return ss.stream.Seek(offset, whence)
+}
+
+func (ss *statsStream) Close() error {
+	atomic.AddUint64(&globalStatistics.Served.TotalBytesServed, uint64(ss.byteCount))
+	ss.byteCount = 0
+
+	return ss.stream.Close()
+}
+
+func (ss *statsStream) Size() int64 {
+	return ss.stream.Size()
+}
+
+func (ss *statsStream) ModTime() time.Time {
+	return ss.stream.ModTime()
 }
 
 func setOtherHeaders(ctx httputils.RequestContext) {
@@ -300,12 +372,15 @@ func getCacheEntryForContentEntry(entry *ContentEntry, originalIdentifier cid.Ci
 		AccessTime: time.Now(),
 	}
 
+	//TODO: make it not require a mutex
 	objectCache[entry.Identifier.String()] = c
 
 	if originalIdentifier.String() != entry.Identifier.String() {
 		objectCache[originalIdentifier.String()] = c
 	}
 
+	atomic.AddUint64(&globalStatistics.Served.TotalHashesServed, 1)
+
 	return c
 }
 
@@ -334,6 +409,7 @@ func IsTrustedPublicKey(key ed25519.PublicKey) bool {
 
 func handle(ctx httputils.RequestContext) {
 	if len(ctx.GetHost()) > 0 && len(ctx.GetTLSServerName()) > 0 && strings.Split(ctx.GetHost(), ":")[0] != ctx.GetTLSServerName() { //Prevents rebinding / DNS stuff
+		atomic.AddUint64(&globalStatistics.HttpCode.Code404, 1)
 		ctx.SetResponseCode(http.StatusNotFound)
 		return
 	}
@@ -350,24 +426,55 @@ func handle(ctx httputils.RequestContext) {
 		}
 		setOtherHeaders(ctx)
 		setCORSHeaders(ctx)
+
+		if ctx.GetPath() == "/stats" {
+			ctx.SetResponseHeader("Content-Type", "application/json")
+			atomic.AddUint64(&globalStatistics.HttpCode.Code200, 1)
+			ctx.SetResponseCode(http.StatusOK)
+
+			statsStruct := struct {
+				Statistics   *statistics `json:"statistics"`
+				TotalEntries uint64      `json:"total_entries"`
+				TotalSize    uint64      `json:"total_size"`
+			}{
+				Statistics: &globalStatistics,
+			}
+
+			if rows, err := statsStatement.Query(); err == nil {
+				defer rows.Close()
+
+				if rows.Next() {
+					rows.Scan(&statsStruct.TotalEntries, &statsStruct.TotalSize)
+				}
+			}
+
+			statBytes, _ := json.MarshalIndent(statsStruct, "", " ")
+			ctx.ServeBytes(statBytes)
+			return
+		}
+
 		pathElements := strings.Split(ctx.GetPath(), "/")
 		if len(pathElements) < 2 {
+			atomic.AddUint64(&globalStatistics.HttpCode.Code400, 1)
 			ctx.SetResponseCode(http.StatusBadRequest)
 			return
 		}
 
 		messageBytes, err := MakyuuIchaival.Bech32Encoding.DecodeString(pathElements[1])
 		if err != nil {
+			atomic.AddUint64(&globalStatistics.HttpCode.Code400, 1)
 			ctx.SetResponseCode(http.StatusBadRequest)
 			return
 		}
 		message := contentmessage.DecodeContentMessage(messageBytes)
 		if message == nil {
+			atomic.AddUint64(&globalStatistics.HttpCode.Code400, 1)
 			ctx.SetResponseCode(http.StatusBadRequest)
 			return
 		}
 
 		if !IsTrustedPublicKey(message.PublicKey) {
+			atomic.AddUint64(&globalStatistics.HttpCode.Code403, 1)
 			ctx.SetResponseCode(http.StatusForbidden)
 			return
 		}
@@ -381,13 +488,16 @@ func handle(ctx httputils.RequestContext) {
 		pTime = cTime
 		cTime = time.Now()
 		if cacheHit {
+			atomic.AddUint64(&globalStatistics.SignatureCache.SignatureCacheHit, 1)
 			ctx.AddTimingInformational("vc", "Ed25519 Cache HIT")
 		} else {
+			atomic.AddUint64(&globalStatistics.SignatureCache.SignatureCacheMiss, 1)
 			ctx.AddTimingInformational("vc", "Ed25519 Cache MISS")
 		}
 		ctx.AddTiming("v", "Ed25519 Verify", cTime.Sub(pTime))
 
 		if !result {
+			atomic.AddUint64(&globalStatistics.HttpCode.Code403, 1)
 			ctx.SetResponseCode(http.StatusForbidden)
 			return
 		}
@@ -400,8 +510,11 @@ func handle(ctx httputils.RequestContext) {
 	} else if ctx.IsOptions() {
 		setOtherHeaders(ctx)
 		setCORSHeaders(ctx)
+
+		atomic.AddUint64(&globalStatistics.HttpCode.Code204, 1)
 		ctx.SetResponseCode(http.StatusNoContent)
 	} else {
+		atomic.AddUint64(&globalStatistics.HttpCode.Code501, 1)
 		ctx.SetResponseCode(http.StatusNotImplemented)
 	}
 }
@@ -518,6 +631,12 @@ func main() {
 	}
 	defer md5Statement.Close()
 
+	statsStatement, err = dbHandle.Prepare("SELECT COUNT(*) AS count, SUM(size) as size FROM entries;")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer statsStatement.Close()
+
 	tlsConfiguration, err := tlsutils.NewTLSConfiguration(*certificatePath, *keypairPath, strings.ToLower(*sniAddressOption))
 	if err != nil {
 		log.Fatal(err)