tinydb/main.go
Přemysl Janouch 7847236dbc
Improve server shutdown
Less likely to corrupt data now.
2019-04-19 12:30:48 +02:00

88 lines
2.0 KiB
Go

// Demos a trivial key-value database backed by a file.
package main
import (
"context"
"io/ioutil"
"log"
"net/http"
"os"
"os/signal"
"strings"
"syscall"
)
func main() {
if len(os.Args) != 3 {
log.Fatalf("Usage: %s LISTEN-ADDRESS DATABASE-FILE\n", os.Args[0])
}
listenAddr, dbFilename := os.Args[1], os.Args[2]
db, err := OpenDB(dbFilename)
if err != nil && os.IsNotExist(err) {
log.Println("database file does not exist, creating")
db, err = CreateDB(dbFilename)
}
if err != nil {
log.Fatalln(err)
}
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
if _, ok := r.URL.Query()["checkpoint"]; ok && r.Method == "GET" {
if err := db.Checkpoint(); err != nil {
w.WriteHeader(http.StatusInternalServerError)
}
return
}
byteValue, err := ioutil.ReadAll(r.Body)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
return
}
key, value := strings.TrimPrefix(r.URL.Path, "/"), string(byteValue)
switch r.Method {
case "GET":
if value, ok, err := db.Get(key); err != nil {
w.WriteHeader(http.StatusInternalServerError)
} else if ok {
w.WriteHeader(http.StatusOK)
w.Write([]byte(value))
} else {
w.WriteHeader(http.StatusNotFound)
}
case "PUT":
if err := db.Put(key, value); err != nil {
w.WriteHeader(http.StatusInternalServerError)
}
case "DELETE":
if err := db.Delete(key); err != nil {
w.WriteHeader(http.StatusInternalServerError)
}
default:
w.WriteHeader(http.StatusMethodNotAllowed)
}
})
server := &http.Server{Addr: listenAddr}
errs := make(chan error, 1)
go func() { errs <- server.ListenAndServe() }()
sig := make(chan os.Signal)
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
select {
case <-sig:
case err := <-errs:
log.Println(err)
}
// For simplicity, we'll wait for everything to finish, including snapshots.
if err := server.Shutdown(context.Background()); err != nil {
log.Fatalln(err)
}
if err := db.Close(); err != nil {
log.Fatalln(err)
}
}