annotation/store/index.go

92 lines
1.9 KiB
Go

package store
import (
"encoding/json"
"fmt"
"hpds_annotation/pkg/utils"
"io"
"os"
"path"
"github.com/klauspost/compress/zstd"
)
func Load(storePath string) map[string]bool {
if !utils.PathExists(storePath) {
_ = os.MkdirAll(storePath, os.ModePerm)
}
fileName := "store"
storeFile := path.Join(storePath, fmt.Sprintf("%s.hdb", fileName))
if !utils.PathExists(storeFile) {
NewFile(storeFile)
return make(map[string]bool)
}
list := make(map[string]bool)
f, _ := os.OpenFile(storeFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
defer func(f *os.File) {
_ = f.Close()
}(f)
buff, err := io.ReadAll(f)
if err != nil {
fmt.Println(err)
return nil
}
if len(buff) > 0 {
str, err := UnCompress(buff)
if err != nil {
return nil
}
err = json.Unmarshal(str, &list)
if err != nil {
return nil
}
}
return list
}
func Save(storePath string, list map[string]bool) {
if !utils.PathExists(storePath) {
_ = os.MkdirAll(storePath, os.ModePerm)
}
fileName := "store"
storeFile := path.Join(storePath, fmt.Sprintf("%s.hdb", fileName))
if !utils.PathExists(storeFile) {
NewFile(storeFile)
}
f, _ := os.OpenFile(storeFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
defer func() {
_ = f.Close()
}()
str, _ := json.Marshal(list)
c := Compress(str)
_, _ = f.Write(c)
}
func NewFile(fileName string) {
f, _ := os.OpenFile(fileName, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
defer func(f *os.File) {
_ = f.Close()
}(f)
}
// Compress 压缩
func Compress(src []byte) []byte {
encoder, _ := zstd.NewWriter(nil)
zstd.WithEncoderConcurrency(3)
return encoder.EncodeAll(src, make([]byte, 0, len(src)))
}
func UnCompress(src []byte) ([]byte, error) {
d, err := zstd.NewReader(nil)
if err != nil {
return nil, err
}
defer d.Close()
uncompressed, err := d.DecodeAll(src, nil)
if err != nil {
return nil, err
}
return uncompressed, nil
}