<commit_before>package main

import (
	"encoding/json"
	"expvar"
	"fmt"
	"net/http"
	"net/http/pprof"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/codahale/hdrhistogram"
)

var expStats *sequinsStats

type sequinsStats struct {
	Qps struct {
		Total    int64
		ByStatus map[string]int64

		total     int64
		status200 int64
		status400 int64
		status404 int64
		status500 int64
		status501 int64
		status502 int64
	}
	Latency struct {
		Max   float64
		Mean  float64
		P9999 float64
		P99   float64
		P95   float64
		P90   float64
		P75   float64
		P50   float64
		P25   float64
	}

	latencyHist *hdrhistogram.Histogram
	queries     chan queryStats

	DiskUsed int64
	lock     sync.RWMutex
}

type queryStats struct {
	duration time.Duration
	status   int
}

func startDebugServer(config sequinsConfig) {
	mux := http.NewServeMux()

	s := &http.Server{
		Addr:    config.Debug.Bind,
		Handler: mux,
	}

	if config.Debug.Expvars {
		mux.HandleFunc("/debug/vars", expvarHandler)
		expStats = newStats(config.LocalStore)
		expvar.Publish("sequins", expStats)
	}

	if config.Debug.Pprof {
		mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
	}

	go s.ListenAndServe()
}

// expvarHandler is copied from the stdlib.
func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	expvar.Do(func(kv expvar.KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

func newStats(localStorePath string) *sequinsStats {
	s := &sequinsStats{
		latencyHist: hdrhistogram.New(0, int64(10*time.Second/time.Microsecond), 5),
		queries:     make(chan queryStats, 1024),
	}

	go s.updateRequestStats()
	go s.updateDiskStats(localStorePath)
	return s
}

func (s *sequinsStats) updateRequestStats() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			s.snapshotRequestStats()

			s.latencyHist.Reset()

			s.Qps.total = 0
			s.Qps.status200 = 0
			s.Qps.status400 = 0
			s.Qps.status404 = 0
			s.Qps.status500 = 0
			s.Qps.status501 = 0
			s.Qps.status502 = 0
		case q := <-s.queries:
			s.latencyHist.RecordValue(int64(q.duration / time.Microsecond))

			s.Qps.total++
			switch q.status {
			case 0, 200:
				s.Qps.status200++
			case 400:
				s.Qps.status400++
			case 404:
				s.Qps.status404++
			case 500:
				s.Qps.status500++
			case 501:
				s.Qps.status501++
			case 502:
				s.Qps.status502++
			default:
				panic("unknown http status code")
			}
		}
	}
}

func (s *sequinsStats) snapshotRequestStats() {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.Qps.Total = s.Qps.total
	s.Qps.ByStatus = make(map[string]int64)
	s.Qps.ByStatus["200"] = s.Qps.status200
	s.Qps.ByStatus["400"] = s.Qps.status400
	s.Qps.ByStatus["404"] = s.Qps.status404
	s.Qps.ByStatus["500"] = s.Qps.status500
	s.Qps.ByStatus["501"] = s.Qps.status501
	s.Qps.ByStatus["502"] = s.Qps.status502

	ms := float64(1000)
	s.Latency.Max = float64(s.latencyHist.Max()) / ms
	s.Latency.Mean = s.latencyHist.Mean() / ms
	s.Latency.P9999 = float64(s.latencyHist.ValueAtQuantile(99.99)) / ms
	s.Latency.P99 = float64(s.latencyHist.ValueAtQuantile(99.0)) / ms
	s.Latency.P95 = float64(s.latencyHist.ValueAtQuantile(95.0)) / ms
	s.Latency.P90 = float64(s.latencyHist.ValueAtQuantile(90.0)) / ms
	s.Latency.P75 = float64(s.latencyHist.ValueAtQuantile(75.0)) / ms
	s.Latency.P50 = float64(s.latencyHist.ValueAtQuantile(50.0)) / ms
	s.Latency.P25 = float64(s.latencyHist.ValueAtQuantile(25.0)) / ms
}

func (s *sequinsStats) updateDiskStats(path string) {
	s.calculateDiskUsage(path)

	ticker := time.NewTicker(1 * time.Minute)
	for range ticker.C {
		s.calculateDiskUsage(path)
	}
}

func (s *sequinsStats) calculateDiskUsage(path string) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			size += info.Size()
		}
		return err
	})

	if err == nil {
		s.lock.Lock()
		defer s.lock.Unlock()

		s.DiskUsed = size
	}
}
<commit_msg>Add a null guard when calculating data size<commit_after>package main

import (
	"encoding/json"
	"expvar"
	"fmt"
	"net/http"
	"net/http/pprof"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/codahale/hdrhistogram"
)

var expStats *sequinsStats

type sequinsStats struct {
	Qps struct {
		Total    int64
		ByStatus map[string]int64

		total     int64
		status200 int64
		status400 int64
		status404 int64
		status500 int64
		status501 int64
		status502 int64
	}
	Latency struct {
		Max   float64
		Mean  float64
		P9999 float64
		P99   float64
		P95   float64
		P90   float64
		P75   float64
		P50   float64
		P25   float64
	}

	latencyHist *hdrhistogram.Histogram
	queries     chan queryStats

	DiskUsed int64
	lock     sync.RWMutex
}

type queryStats struct {
	duration time.Duration
	status   int
}

func startDebugServer(config sequinsConfig) {
	mux := http.NewServeMux()

	s := &http.Server{
		Addr:    config.Debug.Bind,
		Handler: mux,
	}

	if config.Debug.Expvars {
		mux.HandleFunc("/debug/vars", expvarHandler)
		expStats = newStats(config.LocalStore)
		expvar.Publish("sequins", expStats)
	}

	if config.Debug.Pprof {
		mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
	}

	go s.ListenAndServe()
}

// expvarHandler is copied from the stdlib.
func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	expvar.Do(func(kv expvar.KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

func newStats(localStorePath string) *sequinsStats {
	s := &sequinsStats{
		latencyHist: hdrhistogram.New(0, int64(10*time.Second/time.Microsecond), 5),
		queries:     make(chan queryStats, 1024),
	}

	go s.updateRequestStats()
	go s.updateDiskStats(localStorePath)
	return s
}

func (s *sequinsStats) updateRequestStats() {
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			s.snapshotRequestStats()

			s.latencyHist.Reset()

			s.Qps.total = 0
			s.Qps.status200 = 0
			s.Qps.status400 = 0
			s.Qps.status404 = 0
			s.Qps.status500 = 0
			s.Qps.status501 = 0
			s.Qps.status502 = 0
		case q := <-s.queries:
			s.latencyHist.RecordValue(int64(q.duration / time.Microsecond))

			s.Qps.total++
			switch q.status {
			case 0, 200:
				s.Qps.status200++
			case 400:
				s.Qps.status400++
			case 404:
				s.Qps.status404++
			case 500:
				s.Qps.status500++
			case 501:
				s.Qps.status501++
			case 502:
				s.Qps.status502++
			default:
				panic("unknown http status code")
			}
		}
	}
}

func (s *sequinsStats) snapshotRequestStats() {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.Qps.Total = s.Qps.total
	s.Qps.ByStatus = make(map[string]int64)
	s.Qps.ByStatus["200"] = s.Qps.status200
	s.Qps.ByStatus["400"] = s.Qps.status400
	s.Qps.ByStatus["404"] = s.Qps.status404
	s.Qps.ByStatus["500"] = s.Qps.status500
	s.Qps.ByStatus["501"] = s.Qps.status501
	s.Qps.ByStatus["502"] = s.Qps.status502

	ms := float64(1000)
	s.Latency.Max = float64(s.latencyHist.Max()) / ms
	s.Latency.Mean = s.latencyHist.Mean() / ms
	s.Latency.P9999 = float64(s.latencyHist.ValueAtQuantile(99.99)) / ms
	s.Latency.P99 = float64(s.latencyHist.ValueAtQuantile(99.0)) / ms
	s.Latency.P95 = float64(s.latencyHist.ValueAtQuantile(95.0)) / ms
	s.Latency.P90 = float64(s.latencyHist.ValueAtQuantile(90.0)) / ms
	s.Latency.P75 = float64(s.latencyHist.ValueAtQuantile(75.0)) / ms
	s.Latency.P50 = float64(s.latencyHist.ValueAtQuantile(50.0)) / ms
	s.Latency.P25 = float64(s.latencyHist.ValueAtQuantile(25.0)) / ms
}

func (s *sequinsStats) updateDiskStats(path string) {
	s.calculateDiskUsage(path)

	ticker := time.NewTicker(1 * time.Minute)
	for range ticker.C {
		s.calculateDiskUsage(path)
	}
}

func (s *sequinsStats) calculateDiskUsage(path string) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if info != nil && !info.IsDir() {
			size += info.Size()
		}
		return err
	})

	if err == nil {
		s.lock.Lock()
		defer s.lock.Unlock()

		s.DiskUsed = size
	}
}

{'fraction_non_alphanumeric': 0.14466918066506684, 'fraction_numerical': 0.05667923665866758, 'mean_word_length': 3.525336091003102, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1899795', 'n_tokens_mistral': 4117, 'n_tokens_neox': 3543, 'n_words': 639}