Всем привет! 5 февраля завершился очередной VK Cup, в котором в этот раз впервые была секция посвящённая Go. О конкурсе я узнал случайно в одном из Телеграм каналов и решил посмотреть, что же там за задачи. Соревнование состояло из из 3 этапов:
Квалификация: нужно было реализовать несколько функций, чтобы прошли тесты. Дальше проходило 256 человек.
Отбор: задача про внешнюю сортировку и построение кучи, которая не вмещается в RAM. Дальше проходило 16 человек.
Финал: построение коллажа из 1000+ картинок размером 512×512 px.
Про первые 2 раунда я рассказывать в этой статье не буду, возможно сделаю отдельную статью, а расскажу про финал и решение, которое принесло победу. Код решений всех раундов можно посмотреть на GitHub'е.
В финале были чёткий критерий оценки: кто быстрее построит коллаж, тот и победил. Решение «в лоб» решает эту задачу за ~16 секунд на моём AMD Ryzen 7 5800H (16 HT cores). Если интересно как его ускорить до 0.23 секунды, прошу под кат, там много текста, кода, картинок и даже немного ассемблера.
Задача
Вам дан сайт. Чтобы открыть его на порту 8080, скачайте докер образ vkcup2022/golang:latest и запустите его строчкой
docker run -it --rm -d -p 8080:80 --name web vkcup2022/golang:latestНа каждой его странице может быть одно или несколько изображений размером 512 × 512 пикселей в формате PNG, а также одна или несколько ссылок. Обойдите граф сайта в ширину.
В каждой из предложенных картинок замените Red‑ и Green‑составляющие каждого пикселя на 0 таким образом, чтобы осталась только Blue‑составляющая.
После этого объедините все фотографии в единый коллаж именно в том порядке, в котором они встретились вам при обходе в ширину. Ширина коллажа фиксированная: 32 изображения (512 × 32 = 16 384 пикселя). Заполните коллаж слева направо, а потом сверху вниз.
Обратите внимание: на итоговом сайте будет очень большое количество изображений. Но ваше решение должно использовать не более 2 Гбайт оперативной памяти.
Ваша задача — написать как можно более эффективное решение. Чем меньше времени потребуется вашему решению на запуск, тем выше будет ваше место на соревновании. Мы будем запускать ваше решение строкой
./solution url result.pngРешение "в лоб"
package main
import (
"bytes"
"container/list"
"flag"
"fmt"
"image"
"image/color"
"image/png"
"io"
"net/http"
"net/url"
"os"
"regexp"
"time"
)
const (
collageWidth = 32
imageWidth = 512
imageHeight = 512
)
var (
reImg = regexp.MustCompile(`<img src="(/[^"]+.png)"/>`)
reLink = regexp.MustCompile(`<a href="([^"]+)">`)
)
func main() {
flag.Parse()
start := time.Now()
if len(flag.Args()) != 2 {
fatalf("usage %s <url> <result.png>", os.Args[0])
}
addr := flag.Arg(0)
resFile := flag.Arg(1)
fmt.Printf("Start crawling from %s and storing into %s\n", addr, resFile)
u, err := url.Parse(addr)
if err != nil {
fatalf("%v", err)
}
visited := map[string]struct{}{}
queue := list.New()
queue.PushBack(u.Path)
cnt := 0
resImg := image.NewRGBA(image.Rect(0, 0, imageWidth*collageWidth, imageHeight*32))
for queue.Len() > 0 {
location := queue.Front().Value.(string)
queue.Remove(queue.Front())
if _, exists := visited[location]; exists {
continue
}
visited[location] = struct{}{}
links, images, err := getHtmlData(u.Scheme + "://" + u.Host + location)
if err != nil {
fatalf("%v", err)
}
for _, link := range links {
queue.PushBack(link)
}
for _, imgSrc := range images {
img, err := getImage(u.Scheme + "://" + u.Host + imgSrc)
if err != nil {
fatalf("%v", err)
}
for x := 0; x < imageWidth; x++ {
for y := 0; y < imageHeight; y++ {
_, _, b, a := img.At(x, y).RGBA()
resImg.Set((cnt%collageWidth)*imageWidth+x, (cnt/collageWidth)*imageHeight+y, color.RGBA{0, 0, uint8(b), uint8(a)})
}
}
cnt++
}
}
f, err := os.Create(resFile)
if err != nil {
fatalf("%s", err)
}
defer f.Close()
if err := png.Encode(f, resImg); err != nil {
fatalf("%s", err)
}
fmt.Printf("Processed %d file in %s\n", cnt, time.Now().Sub(start))
}
func getHtmlData(url string) ([]string, []string, error) {
resp, err := http.Get(url)
if err != nil {
return nil, nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, nil, fmt.Errorf("http error: %d", resp.StatusCode)
}
buf := bytes.Buffer{}
if _, err := io.Copy(&buf, resp.Body); err != nil {
fatalf("%v", err)
}
var links, images []string
for _, link := range reLink.FindAllStringSubmatch(buf.String(), -1) {
links = append(links, link[1])
}
for _, img := range reImg.FindAllStringSubmatch(buf.String(), -1) {
images = append(images, img[1])
}
return links, images, nil
}
func getImage(url string) (image.Image, error) {
resp, err := http.Get(url)
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("http error: %d", resp.StatusCode)
}
return png.Decode(resp.Body)
}
func fatalf(format string, v ...any) {
_, _ = fmt.Fprintf(os.Stderr, "Error: %s\n", fmt.Sprintf(format, v...))
os.Exit(255)
}Работает долго, почти 16 секунд. Сломается, если картинок будет не 1024 шт., регулярные выражения для поиска ссылок и картинок сломаются, если добавить хотя бы 1 пробел. Надо менять.
Архитектура
Программу можно разделить на 3 блока:
Crawler: HTTP клиент, который умеет:
Скачивать страницы в нужном порядке и находить на них ссылки и картинки.
Скачивать PNG картинки и парсить их.
PNG Copier: основная часть, которая составляет картинки в коллаж. Ядер CPU у нас много, поэтому и Copier'ов должно быть много.
PNG Writer: т.к. подразумевается, что результирующая картинка не влезает в RAM, то нужно уметь записывать выходной файл чанками.
Всё это должно работать максимально параллельно. Общение между частями происходит через каналы. Ниже графическое представление потоков данных:
Далее расскажу как устроена каждая из частей более подробно.
Crawler
Его задача максимально быстро обойти весь сайт и положить в канал все найденные картинки в правильном порядке. Канал я сделал буферизированным на 1024 элемента (количество картинок в тестсете).
Порядок картинок должен быть как при поиске в ширину, значит все найденные ссылки на другие страницы складываем в очередь, а не в стек.
Пока очередь не пуста, вынимаем из головы URL следующей страницы, скачиваем и парсим HTML, найденные ссылки добавляем в хвост списка, URL картинки в канал для Copier'ов. В сокращённом виде выглядит так:
visited := map[string]struct{}{}
queue := list.New()
queue.PushBack(u.Path)
buf := bytes.Buffer{}
for queue.Len() > 0 {
location := queue.Front().Value.(string)
queue.Remove(queue.Front())
if _, exists := visited[location]; exists {
continue
}
visited[location] = struct{}{}
resp, err := c.httpClient.Get(u.Scheme + "://" + u.Host + location)
// ... Обработка ошибок, чтение HTML в buf
links, images, err := parseHTML(buf.String())
if err != nil {
return err
}
for _, link := range links {
queue.PushBack(link)
}
for _, img := range images {
imgCh <- ImgTask{c.id, u.Scheme + "://" + u.Host + img}
}
}Вообще полная версия умеет не ходить за HTML на другие сайты и понимает абсолютные и относительные ссылки. С вероятностью 99,9% в тестовом наборе организаторов этого не было, но на всякий случай реализовал.
Одна из важнейших функций здесь это parseHTML(). Одна из рекомендаций организаторов была не использовать RegExp'ы, а можно ли использовать внешние библиотеки или нет, я, честно говоря, не знаю до сих пор. В любом случае строить DOM нам не надо, а нужно максимально быстро найти нужную информацию, да и велосипеды я люблю. Опыт участия в HighLoad Cup'ах показывает, что с их помощью можно сильно быстрее решить конкретную задачу, когда не нужна универсальность. Мой парсер занимает чуть больше 100 строк, но со своей задачей справляется на отлично, плюс умеет игнорировать ссылки и картинки в комментариях:
func parseHTML(s string) ([]string, []string, error) {
images := make([]string, 0, 1)
links := make([]string, 0, 2)
for len(s) > 0 {
// Find HTML tag start
pos := strings.IndexByte(s, '<')
if pos == -1 {
break
}
s = s[pos+1:]
s = skipSpaces(s)
// Find end of token
pos = strings.IndexAny(s, " \t\n\r>")
if pos == -1 {
return nil, nil, errors.New("invalid html")
}
token := s[:pos]
s = s[pos:]
// Skip comments
if token == "!--" {
pos := strings.Index(s, "-->")
if pos == -1 {
break
}
s = s[pos+3:]
continue
}
// Find end of HTML tag
pos = strings.IndexByte(s, '>')
if pos == -1 {
return nil, nil, errors.New("invalid html")
}
attrs := s[:pos]
s = s[pos+1:]
if attrs != "" {
switch toLower(token) {
case "a":
pos := strings.Index(attrs, "href")
if pos != -1 {
attrs = skipSpaces(attrs[pos+4:])
if attrs[0] == '=' {
attrs = skipSpaces(attrs[1:])
}
if attrs[0] != '"' {
return nil, nil, errors.New("invalid html")
}
pos = strings.IndexByte(attrs[1:], '"')
if pos == -1 {
return nil, nil, errors.New("invalid html")
}
if pos > 0 {
links = append(links, attrs[1:pos+1])
}
}
case "img":
pos := strings.Index(attrs, "src")
if pos != -1 {
attrs = skipSpaces(attrs[pos+3:])
if attrs[0] == '=' {
attrs = skipSpaces(attrs[1:])
}
if attrs[0] != '"' {
return nil, nil, errors.New("invalid html")
}
pos = strings.IndexByte(attrs[1:], '"')
if pos == -1 {
return nil, nil, errors.New("invalid html")
}
if pos > 0 {
images = append(images, attrs[1:pos+1])
}
}
}
}
}
return links, images, nil
}
func skipSpaces(s string) string {
for len(s) > 0 {
switch s[0] {
case ' ', '\t', '\r', '\n':
s = s[1:]
default:
return s
}
}
return s
}
func toLower(s string) string {
b := []byte(s)
for i, c := range b {
if c >= 'A' && c <= 'Z' {
b[i] = c - 'A' + 'a'
}
}
return string(b)
}toLower() захотелось сделать свой, т.к. в стандартном довольно много проверок, а сэкономить несколько лишних наносекунд никогда не помешает. Ниже сравнение:
BenchmarkStdToLower-16 30112750 40.81 ns/op 8 B/op 1 allocs/op
BenchmarkToLower-16 55108370 33.11 ns/op 8 B/op 1 allocs/opВ итоге парсер получился довольно быстрым и практически без выделения памяти в куче, если сравнить с регулярными выражениями, то
BenchmarkRegExp-16 440782 4176 ns/op 723 B/op 10 allocs/op
BenchmarkParser-16 1986980 595.2 ns/op 48 B/op 2 allocs/opМне казалось, что использование пакета bytes вместо strings должно было его ещё ускорить, но оказалось медленнее. К сожалению benchmark не сохранился. Если есть идеи почему, поделитесь в комментариях.
Ещё из интересных моментов в Crawler'е, настройка гошного HTTP клиента:
httpClient: &http.Client{
Transport: &http.Transport{
DisableCompression: true,
MaxIdleConnsPerHost: runtime.NumCPU() * 2,
IdleConnTimeout: time.Minute,
ReadBufferSize: 64 * 1024,
},
},Сжатие не нужно, т.к. HTML маленький, сеть быстрая. Таймаут для ожидающих подключений — минута, чтобы не закрывал раньше времени. И последний штрих — размер буфера для чтения, все HTML и картинки хочется читать за 1 Read.
На самом деле всё это никакой роли не играет, как и распараллеливание обкачки, т.к. страниц относительно мало, вся операция занимает чуть больше 100мс.
PNG Writer
Стандартный encoder умеет записывать только картинку целиком, а это значит придётся разбираться как устроен формат PNG. На сайте LibPNG есть спецификация, описывающая структуру файла и она довольно проста:
Сигнатура
137 80 78 71 13 10 26 10Последовательность чанков, каждый из которых состоит из:
Длина (4 байта)
Тип (4 байта), может быть:
IHDR: Image header (13 байт с основными параметрами)PLTE: PaletteIDAT: Image data (содержимое картинки в сжатом виде с помощью zlib)IEND: Image trailer...
Данные (длинна должна совпадать с числом в первых 4 байтах)x
CRC (4 байта)
IHDR чанк должен быть первым, а IEND - последним. В заголовке содержатся следующие поля:
Width: 4 bytes
Height: 4 bytes
Bit depth: 1 byte
Color type: 1 byte
Compression method: 1 byte (всегда 0)
Filter method: 1 byte
Interlace method: 1 byte
От новой библиотеки для записи PNG требуется функции умеющие:
Записать заголовок (сигнатура +
IHDR)Несколько раз записать данные (
IDAT)Записать признак окончания файла (
IEND)
Так как итоговая высота картинки нам не известна на момент записи заголовка, то сначала записываем любое число в высоту, а в конце нужно будет записать заголовок ещё раз, к счастью он имеет фиксированный размер.
Вообще PNG картинка может иметь много IDAT чанков, т. е. каждый вызов на сохранение контента можно создавать новый чанк, но я решил так не делать, а собрать всю картинку в одном чанке. Для этого был выделен буфер, в который записывается содержимое картинки в сжатом виде. Данные картинки хранятся построчно, перед каждой новой строкой выделен 1 байт для указания фильтра, всего их 5 видов, применяются для лучшего сжатия, но я не использовал никакие. В итоге у меня получился следующий код:
package mypng
import (
"bytes"
"compress/zlib"
"encoding/binary"
"errors"
"hash/crc32"
"image"
"io"
"strconv"
)
const pngHeader = "\x89PNG\r\n\x1a\n"
type Encoder struct {
w io.Writer
header [8]byte
footer [4]byte
tmp [4 * 256]byte
buf *bytes.Buffer
zw *zlib.Writer
}
func NewEncoder(w io.Writer) *Encoder {
buf := bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
zw := zlib.NewWriter(buf)
return &Encoder{
w: w,
buf: buf,
zw: zw,
}
}
func (e *Encoder) WriteIHDR(w, h uint32) error {
if _, err := io.WriteString(e.w, pngHeader); err != nil {
return err
}
binary.BigEndian.PutUint32(e.tmp[0:4], w)
binary.BigEndian.PutUint32(e.tmp[4:8], h)
// Set bit depth and color type.
e.tmp[8] = 8
e.tmp[9] = 6
e.tmp[10] = 0 // default compression method
e.tmp[11] = 0 // default filter method
e.tmp[12] = 0 // non-interlaced
return e.writeChunk(e.tmp[:13], "IHDR")
}
func (e *Encoder) WriteIDAT(i *image.NRGBA) error {
b := i.Bounds()
for y := 0; y < b.Max.Y; y++ {
if _, err := e.zw.Write([]byte{0}); err != nil {
return err
}
offset := (y - b.Min.Y) * i.Stride
if _, err := e.zw.Write(i.Pix[offset : offset+b.Dx()*4]); err != nil {
return err
}
}
return nil
}
func (e *Encoder) FinishIDAT() error {
if err := e.zw.Close(); err != nil {
return err
}
return e.writeChunk(e.buf.Bytes(), "IDAT")
}
func (e *Encoder) WriteIEND() error { return e.writeChunk(nil, "IEND") }
func (e *Encoder) writeChunk(b []byte, name string) error {
n := uint32(len(b))
if int(n) != len(b) {
return errors.New(name + " chunk is too large: " + strconv.Itoa(len(b)))
}
binary.BigEndian.PutUint32(e.header[:4], n)
e.header[4] = name[0]
e.header[5] = name[1]
e.header[6] = name[2]
e.header[7] = name[3]
crc := crc32.NewIEEE()
_, _ = crc.Write(e.header[4:8])
_, _ = crc.Write(b)
binary.BigEndian.PutUint32(e.footer[:4], crc.Sum32())
if _, err := e.w.Write(e.header[:8]); err != nil {
return err
}
if _, err := e.w.Write(b); err != nil {
return err
}
if _, err := e.w.Write(e.footer[:4]); err != nil {
return err
}
return nil
}
Какие‑то части были скопированы из стандартной библиотеки.
Формат картинки NRGBA хранит пикселы в таком же формате, как и PNG с BitDepth=8 и ColorType=6, поэтому можно копировать всю строку целиком, а не попиксельно.
Я написал небольшой benchmark, чтобы посмотреть на узкие места:
package mypng_test
import (
"image"
"os"
"testing"
"round-3/mypng"
)
var testNRGBA *image.NRGBA
func init() {
testNRGBA = image.NewNRGBA(image.Rect(0, 0, 512*32, 512*32))
for i := 0; i < len(testNRGBA.Pix)/4; i++ {
testNRGBA.Pix[i*4+2] = 255
testNRGBA.Pix[i*4+3] = 255
}
}
func BenchmarkWriteRGBA(b *testing.B) {
for i := 0; i < b.N; i++ {
f, err := os.Create("/tmp/test.png")
if err != nil {
b.Fatal(err)
}
defer f.Close()
myEnc := mypng.NewEncoder(f)
if err := myEnc.WriteIHDR(512*32, 512*32); err != nil {
b.Fatal(err)
}
if err := myEnc.WriteIDAT(testNRGBA); err != nil {
b.Fatal(err)
}
if err := myEnc.FinishIDAT(); err != nil {
b.Fatal(err)
}
if err := myEnc.WriteIEND(); err != nil {
b.Fatal(err)
}
}
}После его запуска результаты были такие:
BenchmarkWriteRGBA-16 1 2314487506 ns/op2.3 секунды, беда. Если посмотреть на профайлинг CPU, то видно, что вся нагрузка ложится на zlib.Write() и с этим надо что-то делать.
Идея № 1: zlib имеет разные степени сжатия, одна из них BestSpeed. Кажется то, что надо, пробуем:
BenchmarkWriteRGBA-16 2 770471273 ns/opВ 3 раза лучше, но хочется быстрее. Переписывать zlib времени нет, да и не факт, что получится ускорить (на самом деле можно), значит надо думать как уменьшить объём данных.
Идея № 2: Нам нужен только синий цвет, а сохраняем мы все 4. PNG поддерживает формат с палитрой из 256 цветов, что нам и надо. Каждый оттенок синего однозначно будет соответствовать цвету в палитре. Чтобы сохранять PNG в новом формате, нужно поменять функции записи заголовка и данных, а также добавить новую функцию для сохранения палитры в файле:
func (e *Encoder) WriteIHDR(w, h uint32) error {
if _, err := io.WriteString(e.w, pngHeader); err != nil {
return err
}
binary.BigEndian.PutUint32(e.tmp[0:4], w)
binary.BigEndian.PutUint32(e.tmp[4:8], h)
// Set bit depth and color type.
e.tmp[8] = 8
e.tmp[9] = 3 // !!! 6 -> 3
e.tmp[10] = 0 // default compression method
e.tmp[11] = 0 // default filter method
e.tmp[12] = 0 // non-interlaced
return e.writeChunk(e.tmp[:13], "IHDR")
}
func (e *Encoder) WritePLTE(p color.Palette) error {
if len(p) < 1 || len(p) > 256 {
return errors.New("bad palette length: " + strconv.Itoa(len(p)))
}
for i, c := range p {
c1 := color.NRGBAModel.Convert(c).(color.NRGBA)
e.tmp[3*i+0] = c1.R
e.tmp[3*i+1] = c1.G
e.tmp[3*i+2] = c1.B
}
return e.writeChunk(e.tmp[:3*len(p)], "PLTE")
}
func (e *Encoder) WriteIDAT(i *image.Paletted) error {
b := i.Bounds()
for y := 0; y < b.Max.Y; y++ {
if _, err := e.zw.Write([]byte{0}); err != nil {
return err
}
offset := y * i.Stride
if _, err := e.zw.Write(i.Pix[offset : offset+b.Max.X]); err != nil {
return err
}
}
return nil
}В итоге нужно будет записать в 4 раза меньше данных и это положительно сказалось на скорости:
BenchmarkWriteRGBA-16 2 770471273 ns/op
BenchmarkWritePaletted-16 6 181868631 ns/opВыглядит хорошо, теперь можно собирать всю конструкцию.
Создание коллажа
Копировать последовательно по одной картинке на системе с несколькими ядрами непозволительная роскошь в задачах где цель — скорость, поэтому DataFlow такой:
Есть буфферизованный канал с картинками на одну строку коллажа — чанками (16 384×512 px). Он нужен чтобы не выделять каждый раз память под новую картинку, плюс канал гарантирует, что одновременно в памяти не будет больше чем N картинок, в отличии от
sync.PoolЕсть N воркеров, каждый из которых
Берёт URL картинки и её ID из Crawler'а
Скачивает и декодирует PNG
В случае, если картинка относится к текущему чанку, то копирует её в нужное место чанка, иначе ожидает следующий чанк с помощью SpinLock'а (CPU нам не жалко, а время на переключения контекста жалко).
Если все картинки чанка заполнены, то он передаётся в другой буферизированный канал, откуда другой горутиной он записывается в выходной файл и возвращается в первый канал для повторного использования.
Графически это можно представить так:
А в коде так:
rowImgBufCh := make(chan *image.Paletted, rowImgBufSz)
for i := 0; i < rowImgBufSz; i++ {
rowImgBufCh <- image.NewPaletted(image.Rect(0, 0, imageWidth*collageWidth, imageHeight), nil)
}
rowImgCh := make(chan *image.Paletted, rowImgBufSz)
rowImg := <-rowImgBufCh
line := 0
leftRowImages := int32(collageWidth)
activeWorkers := int32(workers)
for i := 0; i < workers; i++ {
go func() {
for task := range tasksCh {
img, err := c.GetImage(task.Location)
if err != nil {
fatalf("%v", err)
}
// Next line spin
for task.Id/collageWidth > line {
}
copyImage(img, rowImg, (task.Id%collageWidth)*imageWidth)
// Row is complete
if atomic.AddInt32(&leftRowImages, -1) == 0 {
rowImgCh <- rowImg
rowImg = <-rowImgBufCh
leftRowImages = collageWidth
line++
}
}
// All workers are done
if atomic.AddInt32(&activeWorkers, -1) == 0 {
if leftRowImages != collageWidth {
// Clean unused space
for y := 0; y < imageHeight; y++ {
dstPixOffset := rowImg.PixOffset(int((collageWidth-leftRowImages)*imageWidth), y)
for x := 0; x < int(leftRowImages)*imageWidth; x++ {
rowImg.Pix[dstPixOffset+x] = 0
}
}
// Store partial chunk
rowImgCh <- rowImg
line++
}
close(rowImgCh)
}
}()
}
for rowImg := range rowImgCh {
if err := enc.WriteIDAT(rowImg); err != nil {
fatalf("%v", err)
}
rowImgBufCh <- rowImg
}Всё работает без локов, синхронизация идёт с помощью atomic операций. Подготавливая презентацию в последний вечер перед закрытием раунда, я заметил, что мой спинлок довольно сильно выделяется на общем фоне:
После долгих сомнений, стоит ли рисковать и в пятницу ночью переделывать на то, чтобы можно было обрабатывать сразу несколько чанков одновременно, но при этом не нарушить их порядок, я решил убрать спинлок и посмотреть какой будет прирост по скорости. Конечно картинка разъехалась, но выигрыш в скорости был ~10%, а значит надо рискнуть и быстро переделать, хоть после этого уже не будет шанса всё нормально протестировать и исправить.
В итоге активные чанки я сложил в map'у, почему именно в такую структуру данных я и сейчас не отвечу, но описываю как есть. Обернул всё это RWMutex'ом чтобы избежать гонок и паники. Также счётчик незаполненных картинок переехал в информацию о чанке. В итоге код превратился в:
type chunk struct {
rowImg *image.Paletted
leftRowImages int32
}
chunksMap := map[uint16]*chunk{}
chunksMtx := sync.RWMutex{}
line := uint16(0)
activeWorkers := int32(workers)
for i := 0; i < workers; i++ {
go func() {
for task := range tasksCh {
img, err := c.GetImage(task.Location)
if err != nil {
fatalf("%v", err)
}
chunkId := uint16(task.Id / collageWidth)
chunksMtx.Lock()
curChunk := chunksMap[(chunkId)]
if curChunk == nil {
curChunk = &chunk{
rowImg: <-rowImgBufCh,
leftRowImages: int32(collageWidth),
}
chunksMap[chunkId] = curChunk
}
chunksMtx.Unlock()
copyImage(img, curChunk.rowImg, (task.Id%collageWidth)*imageWidth)
// Row is complete
if atomic.AddInt32(&curChunk.leftRowImages, -1) == 0 {
if chunkId != line {
continue
}
// Store chunk
rowImgCh <- curChunk.rowImg
line++
// Check if next chunks have finished earlier
for {
chunksMtx.RLock()
c := chunksMap[line]
chunksMtx.RUnlock()
if c != nil && c.leftRowImages == 0 {
rowImgCh <- c.rowImg
line++
} else {
break
}
}
}
}
// All workers are done
if atomic.AddInt32(&activeWorkers, -1) == 0 {
if lastChunk := chunksMap[line]; lastChunk != nil && lastChunk.leftRowImages != collageWidth {
// Clean unused space
for y := 0; y < imageHeight; y++ {
dstPixOffset := lastChunk.rowImg.PixOffset(int((collageWidth-lastChunk.leftRowImages)*imageWidth), y)
for x := 0; x < int(lastChunk.leftRowImages)*imageWidth; x++ {
lastChunk.rowImg.Pix[dstPixOffset+x] = 0
}
}
// Store partial chunk
rowImgCh <- lastChunk.rowImg
line++
}
close(rowImgCh)
}
}()
}copyImage()
А теперь самая нагруженная часть, копирование сотен миллионов пикселов.
Изначально я решил использовать стандартные методы предоставляемые интерфейсом image.Image для получения цвета пиксела из источника:
func copyImage(src image.Image, dst *image.Paletted, offsetX int) {
for y := 0; y < imageHeight; y++ {
dstPixOffset := dst.PixOffset(offsetX, y)
for x := 0; x < imageWidth; x++ {
_, _, b, _ := src.At(x, y).RGBA()
dst.Pix[dstPixOffset+x] = uint8(b)
}
}
}Если посмотреть на профайлинг
то видно, что бо́льшую часть времени занимает метод At(), плюс он ещё выделяет память. Универсальность здесь не нужна, заранее известен формат тестовых картинок, а именно image.RGBA, поэтому можно избавиться от вызова At() и сразу добраться до нужного пиксела, а точнее его синей составляющей. Я на всякий случай оставил первый вариант, если всё таки придёт не image.RGBA:
func copyImage(src image.Image, dst *image.Paletted, offsetX int) {
switch src := src.(type) {
case *image.RGBA:
for y := 0; y < imageHeight; y++ {
for x := 0; x < imageWidth; x++ {
dst.Pix[dstPixOffset+x] = src.Pix[srcPixOffset+x*4+1]
}
}
default:
for y := 0; y < imageHeight; y++ {
dstPixOffset := dst.PixOffset(offsetX, y)
for x := 0; x < imageWidth; x++ {
_, _, b, _ := src.At(x, y).RGBA()
dst.Pix[dstPixOffset+x] = uint8(b)
}
}
}
}Стало сильно лучше:
BenchmarkCopySTD-16 274 4326141 ns/op 1048586 B/op 262144 allocs/op
BenchmarkCopyPix-16 5012 213639 ns/op 0 B/op 0 allocs/opЕсли посмотреть на то, что происходит на уровне ассемблера, то можно заметить, что копирование одного элемента слайса в другой — это не просто копирование, но и дополнительная проверка на невыход за размеры слайса (Bounds Check).
Любой лишний такт CPU на таких объёмах очень важен, поэтому надо избавляться от неё. Bounds Check можно отключить глобально при компиляции, но управлять этим процессом я не могу, поэтому нужен вариант локального отключения. И здесь понадобится чёрная магия unsafe:
dstAddr := uintptr(unsafe.Pointer(&dst.Pix[0]))
srcAddr := uintptr(unsafe.Pointer(&src.Pix[0]))
for y := 0; y < imageHeight; y++ {
srcPixOffset := src.PixOffset(0, y)
dstPixOffset := dst.PixOffset(offsetX, y)
for x := 0; x < imageWidth; x++ {
//dst.Pix[dstPixOffset+x] = src.Pix[srcPixOffset+x*4+1]
dstPtr := (*uint8)((unsafe.Pointer)(dstAddr + uintptr(dstPixOffset+x)))
srcPtr := (*uint8)((unsafe.Pointer)(srcAddr + uintptr(srcPixOffset+x*4+2)))
*dstPtr = *srcPtr
}
}В строках (1) и (2) достаётся адрес нулевых элементов слайсов, и тут Bounds Check всё ещё существует, но т.к. он вызывается только 1 раз, то не страшно. А дальше в строках (8) и (9) получаются указатели на нужные элементы и в строке (10) происходит фактическое копирование. Это даёт ускорение копирования в 1.5 раза относительно классического копирования элементов слайса. Пруф:
BenchmarkCopySTD-16 274 4326141 ns/op 1048586 B/op 262144 allocs/op
BenchmarkCopyPix-16 5012 213639 ns/op 0 B/op 0 allocs/op
BenchmarkCopyPixUnsafe-16 8310 142169 ns/op 0 B/op 0 allocs/opЕсли теперь посмотреть на ассемблер, то видно как он похудел:
Update: уже после написания этой статьи я нашёл способ как можно ещё ускорить и описал это здесь.
Собственный PNG Decoder
В текущем виде приложение отрабатывает за ~450ms. CPU Profile показывает, что теперь основную часть времени приложение тратит на парсинг входящих PNG файлов, плюс заметная часть времени тратится на создание нового изображения в памяти (выделил стрелкой):
Вообще алгоритм выглядит так:
Прочитать данные PNG в память из сокета
Прочитать сигнатуру и заголовок PNG
Создать новое изображение в памяти
Построчно распаковывать и скопировать данные в изображение созданное в п.3
Проверить CRC PNG чанка
Скопировать изображение из п.3 в чанк коллажа, который будет записан в результирующий файл (
copyImage())
Очевидно, что можно избавиться от лишнего выделения памяти и копирования в пунктах 3–5, а заодно не тратить время на подсчёт CRC, вероятность что организаторы подсунут битый файл стремится к 0.
Писать декодер под все возможные форматы PNG не надо, нужно только для одного типа, но чтобы избежать неприятных сюрпризов я сделал Sniffer — реализацию интерфейса io.Reader, который хранит в себе первые 33 байта, что достаточно для понимания нужный формат PNG пришёл или нет. В случае неизвестного формата, Sniffer можно сбросить в начальное состояние и отдать его стандартному png.Decoder.
package mypng
import (
"io"
)
const bufSize = 33
type Sniffer struct {
r io.Reader
buf [bufSize]byte
left int
}
func NewSniffer(r io.Reader) *Sniffer {
return &Sniffer{
r: r,
left: -1,
}
}
func (s *Sniffer) Read(p []byte) (n int, err error) {
if s.left == -1 {
n, err := io.ReadFull(s.r, s.buf[:])
if err != nil {
return 0, err
}
s.left = n
}
if s.left == 0 {
return s.r.Read(p)
}
if len(p) <= s.left {
copy(p, s.buf[bufSize-s.left:])
s.left -= len(p)
return len(p), err
}
copy(p, s.buf[bufSize-s.left:])
n, err = s.r.Read(p[s.left:])
if err != nil {
return n, err
}
n += s.left
s.left = 0
return n, nil
}
func (s *Sniffer) Restart() {
s.left = bufSize
}Теперь можно заняться своей реализацией декодера, который будет выдавать строку пикселей по запросу. Алгоритм получается следующий:
Получаем данные PNG изображения в память из сокета
Читаем заголовок PNG файла используя Sniffer
Если незнакомый формат, то
Сбрасываем Sniffer как будто из него никто никогда не читал
Возвращаем результат из
png.Decode
Создаём новое изображение типа RawData, которое мимикрирует под интерфейс
image.Imageи умеет возвращать следующую строку пикселов по запросу.В момент копирования внутри функции
copyImage()получаем строку пикселов, в этот момент происходит распаковка необходимого числа байт в статически подготовленный буффер, применяем необходимые PNG фильтры (функция фильтров была честно скопирована из стандартного пакета png) и копируем данные в чанк коллажа.
Таким образом в функции copyImage() появилась ещё 1 ветка в switch:
case *mypng.RawData:
dstAddr := uintptr(unsafe.Pointer(&dst.Pix[0]))
for y := 0; y < imageHeight; y++ {
srcPix := src.Next()
srcAddr := uintptr(unsafe.Pointer(&srcPix[0]))
dstPixOffset := dst.PixOffset(offsetX, y)
for x := 0; x < imageWidth; x++ {
//dst.Pix[dstPixOffset+x] = src.Pix[srcPixOffset+x*4+1]
dstPtr := (*uint8)((unsafe.Pointer)(dstAddr + uintptr(dstPixOffset+x)))
srcPtr := (*uint8)((unsafe.Pointer)(srcAddr + uintptr(x*3+2)))
*dstPtr = *srcPtr
}
}
mypng.RawDataPool.Put(src)Как только объект RawData был скопирован, он отправляется в Pool для повторного использования, что экономит время на лишние выделения памяти.
Такой фокус позволяет сэкономить ещё порядка 130ms (1.4 раза). Код PNG декодера без функции filterPaeth (скопирована из стандартного пакета png) выглядит так:
package mypng
import (
"bytes"
"compress/zlib"
"encoding/binary"
"errors"
"fmt"
"image"
"image/color"
"image/png"
"io"
"sync"
)
var RawDataPool = &sync.Pool{
New: func() any {
return &RawData{
Buffer: bytes.NewBuffer(make([]byte, 0, 8*1024)),
}
},
}
var discardBuf [8 * 1024]byte
type RawData struct {
*bytes.Buffer
cr, pr [512*3 + 1]byte
zr io.ReadCloser
}
func (r *RawData) ColorModel() color.Model { panic("implement me") }
func (r *RawData) Bounds() image.Rectangle { panic("implement me") }
func (r *RawData) At(x, y int) color.Color { panic("implement me") }
func (r *RawData) Reset() {
r.Buffer.Reset()
if r.zr != nil {
r.zr.Close()
r.zr = nil
}
}
func (r *RawData) Next() []byte {
const bytesPerPixel = 3
if r.zr == nil {
r.zr, _ = zlib.NewReader(r.Buffer)
}
// Read the decompressed bytes.
_, err := io.ReadFull(r.zr, r.cr[:])
if err != nil {
panic(err)
}
// Apply the filter.
cdat := r.cr[1:]
pdat := r.pr[1:]
switch r.cr[0] {
case 0:
// No-op.
case 1:
for i := bytesPerPixel; i < len(cdat); i++ {
cdat[i] += cdat[i-bytesPerPixel]
}
case 2:
for i, p := range pdat {
cdat[i] += p
}
case 3:
// The first column has no column to the left of it, so it is a
// special case. We know that the first column exists because we
// check above that width != 0, and so len(cdat) != 0.
for i := 0; i < bytesPerPixel; i++ {
cdat[i] += pdat[i] / 2
}
for i := bytesPerPixel; i < len(cdat); i++ {
cdat[i] += uint8((int(cdat[i-bytesPerPixel]) + int(pdat[i])) / 2)
}
case 4:
filterPaeth(cdat, pdat, bytesPerPixel)
default:
panic(fmt.Errorf("bad filter type %d", r.cr[0]))
}
r.pr, r.cr = r.cr, r.pr
return cdat
}
type pngIHDR struct {
Width uint32
Height uint32
BitDepth uint8
ColorType uint8
CompressionMethod uint8
FilterMethod uint8
InterlaceMethod uint8
}
func Decode(r io.Reader) (image.Image, error) {
sr := NewSniffer(r)
hdr, err := readIHDR(sr)
if err != nil {
return nil, err
}
if hdr.Width != 512 || hdr.Height != 512 {
return nil, errors.New("invalid PNG size")
}
if hdr.BitDepth != 8 || hdr.ColorType != 2 || hdr.CompressionMethod != 0 || hdr.FilterMethod != 0 || hdr.InterlaceMethod != 0 {
sr.Restart()
return png.Decode(sr)
}
chunkId, length, err := readChunkHeader(sr)
if chunkId != "IDAT" {
return nil, errors.New("invalid PNG")
}
rawData := RawDataPool.Get().(*RawData)
rawData.Reset()
if _, err := io.CopyN(rawData, sr, int64(length)); err != nil {
return nil, err
}
// Skip all rest data
for {
if _, err := sr.Read(discardBuf[:]); err != nil {
if err == io.EOF {
break
}
return nil, err
}
}
return rawData, nil
}
func readIHDR(r io.Reader) (pngIHDR, error) {
var tmp [17]byte
// Read signature
_, err := io.ReadFull(r, tmp[:8])
if err != nil {
return pngIHDR{}, err
}
id, length, err := readChunkHeader(r)
if id != "IHDR" || length != 13 {
return pngIHDR{}, errors.New("invalid PNG")
}
// Read header data
if _, err := io.ReadFull(r, tmp[:17]); err != nil {
return pngIHDR{}, err
}
return pngIHDR{
Width: binary.BigEndian.Uint32(tmp[0:4]),
Height: binary.BigEndian.Uint32(tmp[4:8]),
BitDepth: tmp[8],
ColorType: tmp[9],
CompressionMethod: tmp[10],
FilterMethod: tmp[11],
InterlaceMethod: tmp[12],
}, nil
}
func readChunkHeader(r io.Reader) (string, uint32, error) {
var tmp [8]byte
if _, err := io.ReadFull(r, tmp[:]); err != nil {
return "", 0, err
}
return string(tmp[4:8]), binary.BigEndian.Uint32(tmp[:4]), nil
}Сейчас CPU Profile выглядит так:
Здесь видно, что бо́льшая часть времени тратится на применение фильтров и распаковку с помощью zlib. С фильтрами я ничего не придумал, а вот для zlib нашлось решение.
Внешние библиотеки
Как я уже говорил в начале статьи, я до сих пор не знаю можно ли использовать внешние библиотеки или нет, поэтому я отправил 2 решения: с внешними библиотеками и без. Какое из них победило — не знаю ?
Заниматься оптимизацией zlib не было ни времени, ни сил, ни желания. Поэтому я пошёл искать, может кто‑то уже решал данную проблему. И как оказалось есть такой человек, нашлась библиотека github.com/klauspost/compress и она действительно быстрее стандартного пакета zlib. Заодно я решил поменять стандартный http.Client на реализацию из github.com/valyala/fasthttp.
И это принесло ещё почти 20% ускорения. В итоге 16 секунд изначального решения превратились в 0.23 секунды.
Результат работы профайлера итогового решения выглядят так:
Заключение
Спасибо организаторам конкурса и всем кто дочитал этот лонгрид до конца. Если интересна тема оптимизаций приложений, приходите на платформу HighLoad.Fun и вливайтесь в наше комьюнити любителей быстрых "велосипедов". Увидимся на следующем конкурсе.
