image_gen_chromedp.go

package main

import (
	"context"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/chromedp/chromedp"
)

func main() {
	mux := http.NewServeMux()
	fs := http.FileServer(http.Dir("./static"))
	mux.Handle("/static/", http.StripPrefix("/static/", fs))
	srv := &http.Server{
		Addr:    ":8001",
		Handler: mux,
	}

	// サーバはブロックするので別の goroutine で実行する
	go func() {
		if err := srv.ListenAndServe(); err != nil {
			log.Print(err)
		}
	}()

	// シグナルを待つ
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM)
	go takeScreenShot(sigCh)
	<-sigCh

	// シグナルを受け取ったらShutdown
	ctx, _ := context.WithTimeout(context.Background(), 5*time.Second)
	if err := srv.Shutdown(ctx); err != nil {
		log.Print(err)
	}
}

// スクショ撮影
func takeScreenShot(sigCh chan os.Signal) {
	defer func() {
		sigCh <- syscall.SIGTERM
	}()
	ctx, cancel := chromedp.NewContext(context.Background())
	defer cancel()

	// ローカルで起動してるサーバーの特定DOMだけ画像として撮影する
	var buf []byte
	if err := chromedp.Run(ctx, elementScreenshot(`http://localhost:8001/`, `#target`, &buf)); err != nil {
		log.Fatal(err)
	}
	if err := ioutil.WriteFile("result.png", buf, 0644); err != nil {
		log.Fatal(err)
	}
}

// 特定のDOMだけ撮影に使う
func elementScreenshot(urlstr, sel string, res *[]byte) chromedp.Tasks {
	return chromedp.Tasks{
		chromedp.Navigate(urlstr),
		chromedp.Sleep(2 * time.Second),
		chromedp.WaitVisible(sel, chromedp.ByID),
		chromedp.Screenshot(sel, res, chromedp.NodeVisible, chromedp.ByID),
	}
}