wp-metacache/main.go

255 lines
5.3 KiB
Go
Raw Normal View History

2021-08-16 21:35:03 +00:00
package main
import (
2021-08-16 22:10:28 +00:00
"bufio"
2021-08-16 21:35:03 +00:00
"database/sql"
"encoding/json"
"errors"
"flag"
2021-08-16 22:24:33 +00:00
"fmt"
2021-08-16 21:35:03 +00:00
"github.com/elliotchance/phpserialize"
_ "github.com/go-sql-driver/mysql"
"log"
"net"
2021-08-16 23:55:59 +00:00
"net/textproto"
2021-08-16 21:35:03 +00:00
"os"
"sort"
"sync"
"time"
)
// Globals...
var sockpath *string
var dsn *string
var db *sql.DB
type metadataValues struct {
Values []string `php:"v"`
}
type metadataObject struct {
lastHit time.Time
entries map[string]metadataValues
lock sync.RWMutex
loading sync.Once
}
type MetadataCache struct {
objects map[int64]metadataObject
lock sync.RWMutex
objectType string
sizeLimit int
}
type age struct {
key int64
lastHit time.Time
}
func (c *MetadataCache) Purge() {
if len(c.objects) < c.sizeLimit {
// We aren't full so let's not run
return
}
var ages []age
for x := range c.objects {
var a age
a.key = x
a.lastHit = c.objects[x].lastHit
ages = append(ages, a)
}
// Sort by age
sort.Slice(ages, func(a, b int) bool { return ages[a].lastHit.After(ages[b].lastHit) })
// Now we remove all entries from the 90% age mark up
purgeto := int(float64(c.sizeLimit) * 0.9)
for x := purgeto; x < len(ages); x++ {
delete(c.objects, ages[x].key)
}
}
func (c *MetadataCache) Delete(o int64) {
// We could be selective about what we purge but let's see if it matters before we do that
c.lock.Lock()
delete(c.objects, o)
c.lock.Unlock()
}
func (c *MetadataCache) Get(o int64, k string) ([]string, error) {
// Check cache for entry
c.lock.RLock()
object, ok := c.objects[o]
if !ok {
// Object is not in the cache so let's load it up
var tmpobject metadataObject
c.objects[o] = tmpobject
tmpobject.loading.Do(func() {
// Only do this once even if concurrent requests come in
tmpobject.entries = loadDbEntries(c.objectType, o)
})
object = tmpobject
}
2021-08-16 22:19:42 +00:00
c.lock.RUnlock()
2021-08-16 21:35:03 +00:00
object.lock.RLock()
entries, ok := object.entries[k]
2021-08-16 22:15:24 +00:00
object.lock.RUnlock()
2021-08-16 22:19:42 +00:00
2021-08-16 21:35:03 +00:00
object.lock.Lock()
2021-08-16 23:55:59 +00:00
object.lastHit = time.Now()
2021-08-16 21:35:03 +00:00
object.lock.Unlock()
c.lock.Lock()
2021-08-16 23:55:59 +00:00
c.objects[o] = object
2021-08-16 21:35:03 +00:00
c.lock.Unlock()
if !ok {
// Value does not exist! Send it back.
return nil, errors.New("Value not found")
}
return entries.Values, nil
}
func loadDbEntries(ot string, id int64) map[string]metadataValues {
2021-08-16 22:00:04 +00:00
log.Println("loadDbEntries")
2021-08-16 23:55:59 +00:00
var entries map[string]metadataValues
entries = make(map[string]metadataValues)
2021-08-16 21:35:03 +00:00
var table string
var column string
2021-08-16 23:55:59 +00:00
if ot == "u" {
table = "wp_usermeta"
column = "user_id"
} else if ot == "p" {
table = "wp_postmeta"
column = "post_id"
2021-08-16 21:35:03 +00:00
} else {
log.Printf("Invalid object type: %s", ot)
return entries
}
2021-08-16 23:55:59 +00:00
query := fmt.Sprintf("select meta_key, meta_value from %s where %s = ?", table, column)
rows, err := db.Query(query, id)
2021-08-16 21:35:03 +00:00
if err != nil {
log.Printf("db.Query: %s\n", err.Error())
return entries
}
for rows.Next() {
var key string
var value string
rows.Scan(&key, &value)
2021-08-16 23:55:59 +00:00
values, _ := entries[key]
values.Values = append(values.Values, value)
entries[key] = values
}
2021-08-16 21:35:03 +00:00
return entries
}
func init() {
// Get database and socket running
dsn = flag.String("dsn", "", "Database connection string")
sockpath = flag.String("sock", "", "Unix socket path")
2021-08-16 21:45:24 +00:00
flag.Parse()
2021-08-16 22:30:54 +00:00
var err error
db, err = sql.Open("mysql", *dsn)
2021-08-16 21:35:03 +00:00
if err != nil {
log.Fatalf("sql.Open: %s", err.Error())
}
err = db.Ping()
if err != nil {
log.Fatalf("sql.Ping: %s", err.Error())
}
}
func main() {
var UC MetadataCache
UC.objects = make(map[int64]metadataObject)
UC.sizeLimit = 5000
2021-08-16 23:55:59 +00:00
UC.objectType = "u"
2021-08-16 21:35:03 +00:00
var PC MetadataCache
PC.objects = make(map[int64]metadataObject)
PC.sizeLimit = 5000
2021-08-16 23:55:59 +00:00
PC.objectType = "p"
var TC MetadataCache
TC.objects = make(map[int64]metadataObject)
TC.sizeLimit = 5000
TC.objectType = "t"
var CC MetadataCache
CC.objects = make(map[int64]metadataObject)
CC.sizeLimit = 100
CC.objectType = "c"
2021-08-16 21:35:03 +00:00
if err := os.RemoveAll(*sockpath); err != nil {
log.Fatal(err)
}
unixListener, err := net.Listen("unix", *sockpath)
2021-08-16 23:55:59 +00:00
if err != nil {
2021-08-16 21:35:03 +00:00
log.Fatal(err)
}
os.Chmod(*sockpath, 0777)
for {
conn, err := unixListener.Accept()
if err != nil {
// handle error
}
go handleConnection(conn, UC, PC, TC, CC)
2021-08-16 21:35:03 +00:00
}
}
type cachecommand struct {
ObjectType string `json:"t"`
ObjectId int64 `json:"i"`
Key string `json:"k"`
Command string `json:"c"`
}
func handleConnection(conn net.Conn, UC MetadataCache, PC MetadataCache, TC MetadataCache, CC MetadataCache) {
2021-08-16 22:00:04 +00:00
log.Println("handleConnection started")
2021-08-16 21:35:03 +00:00
var m *MetadataCache
2021-08-16 22:10:28 +00:00
reader := bufio.NewReader(conn)
tp := textproto.NewReader(reader)
2021-08-17 00:07:23 +00:00
num:=0
2021-08-16 23:55:59 +00:00
for {
buf, err := tp.ReadLine()
2021-08-16 21:35:03 +00:00
if err != nil {
conn.Close()
2021-08-17 00:14:19 +00:00
break
2021-08-16 21:35:03 +00:00
}
2021-08-17 00:07:23 +00:00
num++
2021-08-16 23:55:59 +00:00
var c cachecommand
err = json.Unmarshal([]byte(buf), &c)
2021-08-16 21:35:03 +00:00
if err != nil {
2021-08-16 23:55:59 +00:00
log.Printf("json.Unmarshal: %s returned %s", buf, err.Error())
}
var values []string
switch c.ObjectType {
case "u":
m=&UC
case "p":
m=&PC
case "t":
m=&TC
case "c":
m=&CC
default:
log.Printf("Invalid cache requested: %s\n", c.ObjectType)
conn.Close()
return
2021-08-16 23:55:59 +00:00
}
if c.Command == "g" {
values, err = m.Get(c.ObjectId, c.Key)
if err != nil {
conn.Write([]byte("404"))
} else {
p, err := phpserialize.Marshal(values, nil)
if err != nil {
log.Fatalf("phpserialize.Marshal: %s", err.Error())
}
conn.Write(p)
}
} else if c.Command == "d" {
m.Delete(c.ObjectId)
conn.Write([]byte("200"))
2021-08-16 21:35:03 +00:00
}
2021-08-17 00:04:20 +00:00
conn.Write([]byte("\n"))
2021-08-16 21:35:03 +00:00
}
2021-08-17 00:07:23 +00:00
log.Printf("Ended connection after %d queries", num)
2021-08-16 21:35:03 +00:00
}