ISUCON スニペット集

他のサーバーにアクセスを転送

二台構成にして、initializeリクエストを別のサーバーに送りたい時

func initialize(c echo.Context) error {
	target, _ := url.Parse("http://192.168.11.2:8080") // 転送先のアプリのURL
	proxy := httputil.NewSingleHostReverseProxy(target)
	req := c.Request().Clone(c.Request().Context())
	req.URL.Path = "/api/initialize2" // 転送先のendpoint
	proxy.ServeHTTP(c.Response(), req)

Memcache

関数と変数定義

var mc *memcache.Client

func getMemcacheKey(keyType string, id interface{}) string {
	return fmt.Sprintf("%s.%v", keyType, id)
}

func fetchIfMemcacheHit[T any](
	cacheKey string,
	onCacheNotHit func() (*T, error),
) (*T, error) {
	item, err := mc.Get(cacheKey)
	if errors.Is(err, memcache.ErrCacheMiss) {
		// cache not hit
		result, err := onCacheNotHit()
		if err != nil {
			return nil, err
		}
		val, err := msgpack.Marshal(result)
		if err != nil {
			return nil, err
		}
		if err := mc.Set(&memcache.Item{Value: val, Key: cacheKey}); err != nil {
			return nil, err
		}
		return result, nil
	} else if err != nil {
		return nil, err
	} else {
		// cache hit
		var result T
		if err := msgpack.Unmarshal(item.Value, &result); err != nil {
			return nil, err
		}
		return &result, nil
	}
}

func fetchIfMemcacheHitBulk[T any](
	cacheKeys []string,
	onCacheNotHit func(idx int) (*T, error),
) ([]T, error) {
	items, err := mc.GetMulti(cacheKeys)
	if err != nil {
		return nil, err
	}
	res := make([]T, len(cacheKeys))
	for i, key := range cacheKeys {
		item := items[key]
		if item == nil {
			// cache not hit
			item, err := onCacheNotHit(i)
			if err != nil {
				return nil, err
			}
			val, err := msgpack.Marshal(item)
			if err != nil {
				return nil, err
			}
			if err := mc.Set(&memcache.Item{Value: val, Key: key}); err != nil {
				return nil, err
			}
			res[i] = *item
			continue
		}

		err := msgpack.Unmarshal(item.Value, &res[i])
		if err != nil {
			return nil, err
		}
	}
	return res, nil
}

初期化

main関数に以下を追記する

	memdAddr := os.Getenv("MEMCACHED_ADDRESS")
	if memdAddr == "" {
		memdAddr = "localhost:11211"
	}
	mc = memcache.New(memdAddr)

initialize関数に以下を追加する

	mc.DeleteAll()

使用例

(はじめに) アップデートのときにキャッシュを消すのを忘れずに

	cacheKey := getMemcacheKey("reaction.count", userID)
	mc.Delete(cacheKey)

sliceの場合

	cacheKey := getMemcacheKey("livestream_tags", livestreamModel.ID)
	res, err := fetchIfMemcacheHit(cacheKey, func() (*[]LivestreamTagModel, error) {
		var res []LivestreamTagModel
		if err := tx.SelectContext(ctx, &res, "SELECT * FROM livestream_tags WHERE livestream_id = ?", livestreamModel.ID); err != nil {
			return nil, err
		}
		return &res, nil
	})
	if err != nil {
		return nil, err
	}

構造体の場合

	cacheKey := getMemcacheKey("user", userId)
	res, err := fetchIfMemcacheHit(cacheKey, func() (*UserModel, error) {
		var res UserModel
		if err := tx.GetContext(ctx, &res, "SELECT * FROM users WHERE id = ?", userId); err != nil {
			return nil, err
		}
		return &res, nil
	})
	if err != nil {
		return nil, err
	}

primitive型の場合

	cacheKey := getMemcacheKey("reaction.count", userID)
	res, err := fetchIfMemcacheHit(cacheKey, func() (*int64, error) {
		var res int64
		if err := tx.GetContext(ctx, &res, "SELECT COUNT(*) FROM reactions where user_id = ?", userID); err != nil {
			return nil, err
		}
		return &res, nil
	})
	if err != nil {
		return nil, err
	}

排他制御

mutex.goなどのファイル名で保存する

package main

import "sync"

type KeyedMutex struct {
	mu    sync.Mutex
	locks map[string]*sync.RWMutex
}

var keyedMutex *KeyedMutex = NewKeyedMutex()

func NewKeyedMutex() *KeyedMutex {
	return &KeyedMutex{
		locks: make(map[string]*sync.RWMutex),
	}
}

func (km *KeyedMutex) getLock(key string) *sync.RWMutex {
	km.mu.Lock()
	defer km.mu.Unlock()

	if lock, exists := km.locks[key]; exists {
		return lock
	}

	km.locks[key] = &sync.RWMutex{}
	return km.locks[key]
}

func (km *KeyedMutex) RLock(key string) {
	lock := km.getLock(key)
	lock.RLock()
}

func (km *KeyedMutex) RUnlock(key string) {
	lock := km.getLock(key)
	lock.RUnlock()
}

func (km *KeyedMutex) RWLock(key string) {
	lock := km.getLock(key)
	lock.Lock()
}

func (km *KeyedMutex) RWUnlock(key string) {
	lock := km.getLock(key)
	lock.Unlock()
}

使い方

以下のようにロックができる

	keyedMutex.RLock("sample key")
	defer keyedMutex.RUnlock("sample key")

ファイルの保存

先に"排他制御"の準備を終わらせる。 file.goなどの名前で保存する。

package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
)

var UPLOADED_FILE_DIR = "/home/isucon/uploaded-files"

func saveFile(content []byte, fileName string) error {
	filePath := filepath.Join(UPLOADED_FILE_DIR, fileName)
	keyedMutex.RWLock(filePath)
	defer keyedMutex.RWUnlock(filePath)

	// remove the old file if it exists
	if err := os.Remove(filePath); err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("failed to delete the old file: %w", err)
		}
	}

	if err := os.WriteFile(filePath, content, os.ModePerm); err != nil {
		return fmt.Errorf("failed to save the file: %w", err)
	}

	return nil
}

// ファイルが存在しない場合は `if errors.Is(err, os.ErrNotExist) {` で判定する
func readFile(fileName string) ([]byte, error) {
	filePath := filepath.Join(UPLOADED_FILE_DIR, fileName)
	keyedMutex.RLock(filePath)
	defer keyedMutex.RUnlock(filePath)

	return os.ReadFile(filePath)
}

initializeで削除するのを忘れずに

	if err := os.RemoveAll(UPLOADED_FILE_DIR); err != nil {
		return echo.NewHTTPError(http.StatusInternalServerError, "failed to remove the directory: "+err.Error())
	}
	if err := os.Mkdir(UPLOADED_FILE_DIR, os.ModeDir|os.ModePerm); err != nil {
		return echo.NewHTTPError(http.StatusInternalServerError, "failed to create the directory: "+err.Error())
	}

idが10000を超過したののみ消したい場合は

func removeExceptInitialImage() error {
	files, err := os.ReadDir(UPLOADED_FILE_DIR)
	if err != nil {
		return err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}
		filename := file.Name()
		userId, err := strconv.Atoi(strings.Split(filename, ".")[0])
		if err != nil {
			return err
		}
		if userId > 10000 {
			err := os.Remove(filepath.Join(UPLOADED_FILE_DIR, filename))
			if err != nil {
				return err
			}
		}
	}
	return nil
}

DBの画像データをストレージに保存する例

main関数の一番下でextractImgFromDB()を呼び出す

func mimeToExtension(mime string) string {
	if mime == "image/jpeg" {
		return ".jpg"
	}
	if mime == "image/png" {
		return ".png"
	}
	if mime == "image/gif" {
		return ".gif"
	}
	return ".unknown"
}

func extractImgFromDB() {
	time.Sleep(10 * time.Second)
	exec.Command("mkdir", "-p", UPLOADED_FILE_DIR).Run()

	rows, err := db.Query("SELECT id, imgdata FROM posts")
	if err != nil {
		log.Fatal(err)
		return
	}
	defer rows.Close()

	for rows.Next() {
		var id int
		var img []byte

		err := rows.Scan(&id, &img)
		if err != nil {
			log.Fatal(err)
			return
		}

		ext := mimeToExtension(http.DetectContentType(img))
		filename := fmt.Sprintf("%d%s", id, ext)
		err = saveFile(img, filename)
		if err != nil {
			log.Fatal(err)
			return
		}
		fmt.Println(filename, "saved")
	}

	if err = rows.Err(); err != nil {
		log.Fatal(err)
	}
}