Initial commit

web/api.go

@@ -0,0 +1,120 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/go-redis/redis/v7"
+	"github.com/mailru/easyjson"
+	. "github.com/simon987/imhashdb"
+	"log"
+	"os"
+	"runtime/pprof"
+	"time"
+)
+
+func submitQuery(value string) bool {
+	if Rdb.SIsMember(wipQueue, value).Val() {
+		return false
+	}
+
+	if Rdb.Exists(outQueue+value).Val() == 1 {
+		return false
+	}
+
+	Rdb.ZAdd(inQueue, &redis.Z{
+		Score:  float64(time.Now().Unix()),
+		Member: value,
+	})
+	return true
+}
+
+func pollQuery(ctx context.Context, reqStr string) ([]byte, error) {
+	key := outQueue + reqStr
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, errors.New("timeout")
+		default:
+		}
+
+		value, _ := Rdb.Get(key).Bytes()
+		if value != nil {
+			return value, nil
+		}
+		time.Sleep(time.Millisecond * 50)
+	}
+}
+
+func query(c *gin.Context) {
+	var req QueryReq
+	err := c.BindJSON(&req)
+	if err != nil {
+		c.JSON(400, gin.H{"err": "Invalid request"})
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*20)
+	defer cancel()
+
+	reqJson, _ := easyjson.Marshal(req)
+	value := string(reqJson)
+
+	submitQuery(value)
+	b, err := pollQuery(ctx, value)
+	if err != nil {
+		b, _ = easyjson.Marshal(QueryResp{
+			Err: err.Error(),
+		})
+	}
+	c.Data(200, gin.MIMEJSON, b)
+}
+
+func hash(c *gin.Context) {
+	var req HashReq
+	err := c.BindJSON(&req)
+	if err != nil {
+		c.JSON(400, gin.H{"err": "Invalid request"})
+		return
+	}
+
+	h, err := ComputeHash(req.Data)
+	if err != nil {
+		c.JSON(500, gin.H{"err": "Couldn't compute image hash"})
+		return
+	}
+
+	b, _ := easyjson.Marshal(HashResp{
+		AHash: h.AHash.Bytes, DHash: h.DHash.Bytes,
+		MHash: h.MHash.Bytes, PHash: h.PHash.Bytes,
+		WHash: h.WHash.Bytes,
+	})
+	c.Data(200, gin.MIMEJSON, b)
+}
+
+func main() {
+	Init()
+
+	f, err := os.Create("prof")
+	if err != nil {
+		log.Fatal(err)
+	}
+	pprof.StartCPUProfile(f)
+	go func() {
+		time.Sleep(time.Second * 15)
+		pprof.StopCPUProfile()
+		fmt.Println("!!!!!!!!!!!!!!!")
+		f.Close()
+	}()
+
+	r := gin.Default()
+	r.POST("/api/hash", hash)
+	r.POST("/api/query", query)
+
+	//TODO: concurrency
+	go queryWorker()
+
+	r.Run()
+}

web/worker.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"context"
+	"github.com/mailru/easyjson"
+	. "github.com/simon987/imhashdb"
+	"go.uber.org/zap"
+	"time"
+)
+
+const inQueue = "qq:in"
+const outQueue = "qq:out:"
+const wipQueue = "qq:wip"
+
+const CacheLength = time.Second * 30
+
+func queryWorker() {
+	Logger.Info("Query worker started")
+	for {
+		value := Rdb.BZPopMin(time.Second * 30, inQueue).Val()
+		if value == nil {
+			continue
+		}
+		Logger.Info("worker query start")
+		member := value.Member.(string)
+		var req QueryReq
+		_ = easyjson.Unmarshal([]byte(member), &req)
+
+		resp, err := dbQuery(req, member)
+
+		var b []byte
+		if err != nil {
+			Logger.Warn("worker query error", zap.Error(err))
+			b, _ = easyjson.Marshal(QueryResp{
+				Err: err.Error(),
+			})
+		} else {
+			Logger.Info("worker query done")
+			b = resp
+		}
+		Rdb.Set(outQueue + member, b, CacheLength)
+	}
+}
+
+func dbQuery(req QueryReq, value string) ([]byte, error) {
+	Rdb.SAdd(wipQueue, value)
+	Rdb.Expire(wipQueue, time.Minute * 10)
+
+	defer Rdb.SRem(wipQueue, value)
+
+	resp, err := FindImagesByHash(context.Background(), req.Hash, req.HashType, req.Distance, req.Limit, req.Offset)
+	if err != nil {
+		Logger.Error("Couldn't perform query")
+		return nil, err
+	}
+
+	return resp, nil
+}