feat: improvie update mechanism (#14)

.gitignore

@@ -8,3 +8,4 @@ cov.out
 # massive data files
 /pwned-passwords-*
 *.bin
+pwned-passwords.lock

cmd/pwnd/main.go

@@ -10,12 +10,16 @@ import (
 func main() {
 
 	// parse flags
-	var dbFile string
+	var (
+		dbFile        string
+		updatedDbFile string
+	)
 	flag.StringVar(&dbFile, "database", pwnedpass.DatabaseFilename, "path to the database file")
+	flag.StringVar(&updatedDbFile, "updated-database", pwnedpass.UpdatedDatabaseFilename, "path to the database file")
 	flag.Parse()
 
 	// open the offline database
-	od, err := pwnedpass.NewOfflineDatabase(dbFile)
+	od, err := pwnedpass.NewOfflineDatabase(dbFile, updatedDbFile)
 	if err != nil {
 		panic(err)
 	}

cmd/pwngen/main.go

@@ -25,7 +25,8 @@ const IndexSegmentSize = 256 << 16 << 3 // exactly 256^3 MB
 
 // DatabaseFilename indicates the default location of the database file
 // to be created.
-var DatabaseFilename = "pwned-passwords.bin"
+var DatabaseFilename = "updated-pwned-passwords.bin"
+var LockFileName = "pwned-passwords.lock"
 
 type loadWorker struct {
 	sugar *zap.SugaredLogger
@@ -40,6 +41,7 @@ type response struct {
 // The work that needs to be performed
 // The input type should implement the WorkFunction interface
 func (w loadWorker) Run(ctx context.Context) interface{} {
+	os.Create(LockFileName)
 	var body []byte
 	for i := 0; i < 3; i++ {
 		httpClient := &http.Client{}
@@ -192,5 +194,6 @@ func main() {
 	}
 
 	sugar.Infof("OK")
+	os.Remove(LockFileName)
 
 }

go.mod

@@ -7,6 +7,8 @@ require (
 	golang.org/x/exp v0.0.0-20230905200255-921286631fa9
 )
 
+require github.com/robfig/cron/v3 v3.0.0 // indirect
+
 require (
 	github.com/tejzpr/ordered-concurrently/v3 v3.0.1
 	go.uber.org/multierr v1.10.0 // indirect

go.sum

@@ -1,6 +1,8 @@
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
+github.com/robfig/cron/v3 v3.0.0 h1:kQ6Cb7aHOHTSzNVNEhmp8EcWKLb4CbiMW9h9VyIhO4E=
+github.com/robfig/cron/v3 v3.0.0/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
 github.com/tejzpr/ordered-concurrently/v3 v3.0.1 h1:TLHtzlQEDshbmGveS8S+hxLw4s5u67aoJw5LLf+X2xY=
 github.com/tejzpr/ordered-concurrently/v3 v3.0.1/go.mod h1:mu/neZ6AGXm5jdPc7PEgViYK3rkYNPvVCEm15Cx/iRI=
 go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=

offline.go

@@ -8,20 +8,25 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"os"
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
+	"github.com/robfig/cron/v3"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"golang.org/x/exp/mmap"
 )
 
 const (
 	// DatabaseFilename is the default path to the database.
-	DatabaseFilename = "pwned-passwords.bin"
+	DatabaseFilename        = "pwned-passwords.bin"
+	UpdatedDatabaseFilename = "updated-pwned-passwords.bin"
+	LockFileName            = "pwned-passwords.lock"
 
 	// IndexSegmentSize is the exact size of the index segment in bytes.
 	IndexSegmentSize = 256 << 16 << 3 // exactly 256^3 MB
@@ -55,6 +60,7 @@ type (
 	OfflineDatabase struct {
 		database readCloserAt
 		logger   zap.Logger
+		cron     *cron.Cron
 	}
 
 	// readCloserAt is an io.ReaderAt that can be Closed and whose
@@ -71,13 +77,8 @@ type (
 
 // NewOfflineDatabase opens a new OfflineDatabase using the data in the given
 // database file.
-func NewOfflineDatabase(dbFile string) (*OfflineDatabase, error) {
-
-	db, err := mmap.Open(dbFile)
-	if err != nil {
-		return nil, fmt.Errorf("error opening index: %s", err)
-	}
-
+func NewOfflineDatabase(dbFile string, updatedDbFile string) (*OfflineDatabase, error) {
+	lockExists := false
 	encoderCfg := zap.NewProductionEncoderConfig()
 	encoderCfg.TimeKey = "timestamp"
 	encoderCfg.EncodeTime = zapcore.ISO8601TimeEncoder
@@ -99,17 +100,68 @@ func NewOfflineDatabase(dbFile string) (*OfflineDatabase, error) {
 			"pid": os.Getpid(),
 		},
 	}
+	logger := *zap.Must(config.Build())
+	for {
+		if _, err := os.Stat(LockFileName); err == nil {
+			lockExists = true
+		}
+		if _, err := os.Stat(dbFile); err == nil {
+			break
+		} else {
+			// Check if error indicates a missing file?
+			if os.IsNotExist(err) && lockExists {
+				logger.Warn("Lock file exists, but database file does not. Waiting for lock to be released.")
+				time.Sleep(1 * time.Minute)
+			}
+		}
+	}
+
+	if _, err := os.Stat(updatedDbFile); err == nil {
+		if !lockExists {
+			if err := os.Rename(updatedDbFile, dbFile); err != nil {
+				return nil, fmt.Errorf("error moving updated database: %s", err)
+			}
+		}
+	}
+
+	db, err := mmap.Open(dbFile)
+	if err != nil {
+		return nil, fmt.Errorf("error opening index: %s", err)
+	}
+	c := cron.New()
 	odb := &OfflineDatabase{
 		database: db,
-		logger:   *zap.Must(config.Build()),
+		logger:   logger,
+		cron:     c,
 	}
 
+	c.AddFunc("@hourly", func() {
+		if _, err := os.Stat(updatedDbFile); err == nil {
+			lockExists := false
+			if _, err := os.Stat(LockFileName); err == nil {
+				lockExists = true
+			}
+			if !lockExists {
+				db.Close()
+				if err := os.Rename(updatedDbFile, dbFile); err != nil {
+					log.Panic(err)
+				}
+				db, err := mmap.Open(dbFile)
+				if err != nil {
+					log.Panic(err)
+				}
+				odb.database = db
+			}
+		}
+	})
+	c.Start()
 	return odb, nil
 
 }
 
 // Close frees resources associated with the database.
 func (od *OfflineDatabase) Close() error {
+	od.cron.Stop()
 	return od.database.Close()
 }
 

offline_test.go

@@ -25,7 +25,7 @@ func TestPwned(t *testing.T) {
 		},
 	}
 
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		t.Fatalf("unexpected error: %s", err)
 	}
@@ -44,7 +44,7 @@ func TestPwned(t *testing.T) {
 
 func BenchmarkPwned(b *testing.B) {
 
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		b.Fatalf("unexpected error: %s", err)
 	}
@@ -93,7 +93,7 @@ func TestScan(t *testing.T) {
 		},
 	}
 
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		t.Fatalf("unexpected error: %s", err)
 	}
@@ -127,7 +127,7 @@ func BenchmarkScan(b *testing.B) {
 		EndPrefix   = [3]byte{0x05, 0x31, 0x91}
 	)
 
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		b.Fatalf("unexpected error: %s", err)
 	}
@@ -179,7 +179,7 @@ func TestLookup(t *testing.T) {
 		},
 	}
 
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		t.Fatalf("unexpected error: %s", err)
 	}
@@ -206,7 +206,7 @@ func TestLookup(t *testing.T) {
 func BenchmarkHTTPPassword(b *testing.B) {
 
 	// open the offline database
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		panic(err)
 	}
@@ -242,7 +242,7 @@ func BenchmarkHTTPPassword(b *testing.B) {
 func BenchmarkHTTPRange(b *testing.B) {
 
 	// open the offline database
-	od, err := NewOfflineDatabase(DatabaseFilename)
+	od, err := NewOfflineDatabase(DatabaseFilename, UpdatedDatabaseFilename)
 	if err != nil {
 		panic(err)
 	}