Пакет context в Go: взгляд профессионала

Написано в соавторстве с Константином Соколовым, это бэкап нашего поста c Habrahabr.

Обзор

Пакет context появился как x/net/context в 2014 году и быстро обрел популярность. В 2016 году его добавили в стандартную библиотеку Go 1.7. С тех пор практически ни одно приложение на Go не обходится без его использования, потому что пакет ощутимо упрощает многие задачи.

Что такое контекст

Контекстом называют интерфейс Context из пакета context.

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key any) any
}

Код почти дословно иллюстрирует, для чего используется контекст:

• чтобы устанавливать дедлайн исполнения блока кода;
• оповещать об окончании исполнения блока кода;
• узнавать причину отмены контекста;
• получать значения по ключу.

Интерфейса достаточно для использования в любых местах, где код может «зависнуть». Это любое сетевое взаимодействие, а также долгие задачи, не выходящие за рамки процесса ОС. Кроме того, контекст можно использовать для неявной передачи параметров в функции.

Все методы context.Context могут вызываться одновременно из нескольких горутин.

В Go принято передавать контекст в виде первого аргумента функции. Обычно его называют ctx.

Вот пример функции из пакета net/http:

func NewRequestWithContext(ctx context.Context, method, url string, body io.Reader) (*Request, error) {
	...
}

Корневой контекст

Корневой контекст невозможно отменить, он не имеет дедлайна и не содержит значений.

Обычно корневой контекст создается в main() приложения и оборачивается в производный контекст с отменой, который затем передается в функции. Когда контекст отменяется в main(), отменяется и код в глубине стека вызовов. Это помогает реализовать корректное завершение работы процесса.

Есть две функции получения корневого контекста:

• context.Background()
• context.TODO()

Реализованы они очень просто:

func Background() Context {
    return backgroundCtx{}
}

func TODO() Context {
    return todoCtx{}
}

...

type backgroundCtx struct{ emptyCtx }

type todoCtx struct{ emptyCtx }

Оба контекста по сути ничем не отличаются, но context.TODO() используется только на время проектирования или эксперимента. context.TODO() пригождается, когда нужно вызвать функцию, требующую context.Context в качестве аргумента, но в данный момент контекст не проброшен сверху по стеку вызовов. Тогда нормально временно использовать context.TODO(), а позже пробросить полноценный контекст, который можно отменять.

Создание нового контекста

Все контексты, за исключением корневых, создаются на основе родительских.

Новый контекст расширяет родительский новой информацией. Например, в новый контекст можно положить ключ со значением, при этом не модифицируя родительский контекст.

Для создания новых контекстов в пакете предусмотрены функции WithCancel, WithDeadline, WithTimeout, WithValue и WithoutCancel.

WithCancel

ctx, cancel := context.WithCancel(context.Background())
defer cancel() // cancel when we are finished

WithCancel возвращает новый контекст и функцию отмены. Новый контекст хранит ссылку на родительский контекст. Функцию отмены нужно вызывать на том же уровне, где она была создана. Это называется отменой вручную. Вызывать функцию отмены в другом месте в Go считается антипаттерном, поскольку это может приводить к утечке незакрытых контекстов.

Вызов cancel закрывает канал, возвращаемый методом контекста Done(). Также канал Done() дочернего контекста закрывается, если закрывается канал Done() родительского контекста.

WithDeadline и WithTimeout

Автоматическую отмену контекста по времени можно сделать при помощи функций WithDeadline и WithTimeout. Они позволяют ограничить максимальное время выполнения блока кода.

Пример:

func slowOperationWithTimeout(ctx context.Context) (Result, error) {
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()
	return slowOperation(ctx)
}

Здесь функция context.WithTimeout принимает два аргумента: существующий контекст и time.Duration. Контекст будет отменен по истечении таймаута автоматически, но также обязательно отменять контекст вручную при помощи defer, чтобы не допустить утечки ресурсов (памяти и горутин). Допустимо делать отмену контекста несколько раз. После первого вызова все последующие игнорируются.

Функция context.WithDeadline отличается тем, что вторым аргументом принимает time.Time. Если указать уже прошедший момент времени, то будет создан сразу отмененный контекст. Метод Deadline экземпляра типа context.Context помогает узнать, в какой момент произойдет отмена. Метод возвращает time.Time и ok. Deadline возвращает ok=false, когда дедлайн не установлен.

WithValue

В тех случаях, когда невозможно передать данные явным образом через несколько промежуточных функций, можно использовать WithValue. WithValue добавляет пару «ключ — значение» в контекст. Функция принимает три аргумента: родительский контекст, ключ и значение. Во избежание коллизий ключ не должен быть одним из встроенных типов и должен поддерживать операторы сравнения == и !=. 

type favContextKey string

key := favContextKey("language")
ctx := context.WithValue(context.Background(), key, "Go")

Получить значение можно методом Value. Он принимает ключ и возвращает его значение. Если ключ не был найден, возвращается nil.

Пример:

ctxValue := ctx.Value(key)

Обработка ошибок

Узнавать об отмене контекста помогают методы Done и Err экземпляра context.Context.

Done представляет канал типа struct{}, который закрывается при отмене контекста и возвращает нулевое значение (nil) при последующей попытке чтения из него.

Метод Err может возвращать три варианта значений:

• nil — если контекст активен;
• context.Canceled типа error в случае явной отмены контекста;
• context.DeadlineExceeded в случае отмены по истечении времени.

Новое в 1.21

WithoutCancel

В Go 1.21 добавлена функция WithoutCancel.

WithoutCancel возвращает копию родительского контекста, которая не будет отменена при отмене родительского контекста.

WithoutCancel можно использовать в следующих довольно частых сценариях:

• при обработке операций rollback/cleanup в контексте какого-либо события (например, обработки HTTP-запроса). Операции должны продолжиться, несмотря на отмену самого события (например, клиент API ушел или наступил таймаут операции)
• при обработке длительных операций, запущенных каким-либо событием (например, HTTP-запросом). Событие можно пометить как обработанное, не дожидаясь завершения длительной операции.

Полученный контекст не возвращает Deadline или Err. Значение канала Done равно nil. Поэтому чтение из него приведет к блокировке программы навсегда.

func WithoutCancel(parent Context) Context {
	if parent == nil {
		panic("cannot create context from nil parent")
	}
	return withoutCancelCtx{parent}
}

type withoutCancelCtx struct {
	c Context
}

func (withoutCancelCtx) Deadline() (deadline time.Time, ok bool) {
	return
}

func (withoutCancelCtx) Done() <-chan struct{} {
	return nil
}

func (withoutCancelCtx) Err() error {
	return nil
}

AfterFunc

В Go 1.21 добавлена новая функция AfterFunc.

Иногда нужно отменить блокирующую функцию, которая поддерживает прерывание своей работы, но не через механизм отмены контекста. До Go 1.21 это делалось сложно и неэффективно. Например, когда нужно было прервать чтение или запись в net.Conn или ожидание на sync.Cond, это делалось через запуск отдельной горутины. Горутина ждала, когда отменится контекст, а затем прерывала блокирующую функцию. Запуск новых горутин достаточно эффективен, но это большие накладные расходы, которые не имеют смысла, если операция очень легкая и отрабатывает за миллисекунды.

AfterFunc позволяет зарегистрировать функцию, которая вызывается при отмене контекста.

func AfterFunc(ctx Context, f func()) (stop func() bool)

AfterFunc принимает функцию f, которая будет выполнена после того, как контекст будет отменен (в том числе после истечения таймаута). Если контекст уже завершен, функция будет запущена сразу же. Выполнение f происходит в отдельной горутине. Множественные вызовы AfterFunc выполняются независимо друг от друга.

AfterFunc возвращает стоп-функцию. Вызов стоп-функции разрывает ассоциацию f с контекстом. Стоп-функция возвращает false, если контекст уже в состоянии Done и f уже была запущена, или если f уже была остановлена. Стоп-функция возвращает true, если f была прервана. Стоп-функция не дожидается, пока f будет завершена, поэтому если нужно контролировать состояние, то лучше явно коммуницировать с f.

Оптимизация инициализации пакета

Рантайм Go с каждой версией становится лучше. Иногда для оптимизации рантайма требуются изменения в стандартной библиотеке.

Пакет context эти оптимизации не обошли стороной. До Go 1.21 функции Background и TODO возвращали заранее заготовленные экземпляры контекста.

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context {
	return background
}

В версии 1.21 контексты Background и TODO стали разными типами, встраивающими пустой контекст. То есть теперь нет глобальных переменных и инициализация рантайма проходит чуть-чуть быстрее.

type backgroundCtx struct{ emptyCtx }

type todoCtx struct{ emptyCtx }

...

func Background() Context {
    return backgroundCtx{}
}

Особенности использования контекста

Do not store Contexts inside a struct type.

Документация пакета context настоятельно рекомендует не хранить контекст внутри структур. Вместо этого нужно явно передавать контекст в функции отдельным аргументом.

Хранение контекста в структуре ведет к путанице и непредсказуемому поведению. Например, если в структуре несколько методов используют контекст, то невозможно выставить дедлайн или выполнить отмену только для одного из методов. Код, использующий такую структуру, сложнее отлаживать и тестировать в первую очередь потому, что для аккуратного использования методов структуры разработчику нужно будет прочитать все реализации методов структуры. А то, что нужно погружаться в реализацию при использовании каких-то методов может означать только одно: при проектировании нарушены базовые принципы абстракции и инкапсуляции.

При явной передаче контекста в функции разработчики могут в полной мере воспользоваться его расширяемостью, устанавливать для каждого вызова дедлайн, отмену и значения ключей.

Передача значений контекста

Use context Values only for request-scoped data that transits processes and APIs, not for passing optional parameters to functions.

Одно из преимуществ использования WithValue — это возможность неявной передачи данных через несколько промежуточных функций. Однако эту функциональность стоит использовать взвешенно. И вот почему. 

Если передавать в контексте опциональные аргументы функций, это усложняет их документирование. В комментарии к функции будет некорректно писать, что функция ожидает в контексте определенный ключ. Такие вещи должны проверяться компилятором, а не программистом при написании кода. Для опциональных аргументов лучше использовать обычные аргументы функций с нулевыми значениями.

В контексте правильно передавать только значения, которые пересекают границы процессов и «живут» от начала и до конца обработки задачи. Например, идентификатор пользователя или запроса хорошо ложится в контекст. В любой функции, использующейся по ходу обработки запроса, идентификатор можно залогировать или передать во внешнюю систему. Логгер тоже хорошо ложится в контекст. Это особенно актуально для структурированного логирования.

Особенности реализации контекста

В пакете Context в версии 1.21 есть следующие типы, реализующие интерфейс Context:

• emptyCtx
• backgroundCtx
• todoCtx
• cancelCtx
• timerCtx
• afterFuncCtx
• valueCtx
• withoutCancelCtx
• stopCtx

Embedding

Пакет context активно использует механизм встраивания (embedding) в качестве альтернативы традиционному механизму наследования.

Например, backgroundCtx реализован как type backgroundCtx struct{ emptyCtx }

Тип emptyCtx встроен в тип backgroundCtx.

У emptyCtx реализованы методы Deadline(), Done(), Err(), Value(), а backgroundCtx реализует лишь свой метод String. При этом все публичные методы emptyCtx автоматически доступны при обращении к backgroundCtx.

Реализация AfterFunc

Рассмотрим реализацию AfterFunc:

func AfterFunc(ctx Context, f func()) (stop func() bool) {
	a := &afterFuncCtx{
		f: f,
	}
	a.cancelCtx.propagateCancel(ctx, a)
	return func() bool {
		stopped := false
		a.once.Do(func() {
			stopped = true
		})
		if stopped {
			a.cancel(true, Canceled, nil)
		}
		return stopped
	}
}

...

type afterFuncCtx struct {
	cancelCtx
	once sync.Once
	f    func()
}

AfterFuncCtx хранит объект sync.Once.

a.once.Do() вызывает func() { stopped = true } только в том случае, если она вызывается в первый раз. Do() потокобезопасна, то есть Do() можно одновременно вызвать из нескольких горутин. Только одна из горутин выполнит f. Так сделано для того, чтобы при отмене контекста из разных горутин AfterFunc выполнилась только один раз.

Reusable closed channel

В пакете context есть интересная глобальная переменная closedchan.

var closedchan = make(chan struct{})

func init() {
	close(closedchan)
}

Этот код создает глобальный закрытый канал на старте приложения, импортирующего пакет context. Далее закрытый канал используется для изменения и проверки состояния Done при отмене контекста.

Для начала ознакомимся с реализацией cancelCtx (такую структуру создает WithCancel):

type cancelCtx struct {
	Context               // Родительский контекст
	done     atomic.Value // Хранит chan struct{}, канал закрывается на первом вызове cancel()
	... // Остальные поля
}

Вот что происходит при вызове cancelCtx.Done() (схематично, в реализации еще есть мьютекс):

func (c *cancelCtx) Done() <-chan struct{} {
	d := c.done.Load() // Загружаем значение атомарной переменной
	if d == nil { // Если значение не установлено, создаем новый канал и сохраняем его
        	d = make(chan struct{})
        	c.done.Store(d)
	}

	return d.(chan struct{})
}

При вызове cancel() используем closedchan:

func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
	...
	d, _ := c.done.Load().(chan struct{}) // Загружаем содержимое атомарной переменной
	if d == nil {
		c.done.Store(closedchan) // Если Done() еще никто не вызывал, записываем заранее заготовленный закрытый канал
	} else {
		close(d) // Если Done() вызывали, канал был создан, нужно его закрыть
	}
	...
}

Также используем closedchan для проверки того, что родительский контекст закрыт:

func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	done := parent.Done()
	if done == closedchan || done == nil { // Удобно, потому что для проверки канала на закрытость не нужно его читать
		return nil, false
	}
}

Такая реализация позволяет не создавать лишние каналы и немного оптимизирует отмену контекстов.

Что происходит при отмене контекста, созданного WithCancel

Каждый контекст связан со своим родительским контекстом, поэтому при отмене родительского контекста все дочерние будут также отменены.

Функция propagateCancel отменяет дочерние контексты при отмене родительского. Она вызывается в функции отмены контекста вручную (withCancel), отмены по дедлайну (WithDeadline) и AfterFunc. Если родительский контекст уже отменен, то она запускает отмену всех дочерних контекстов при помощи функции cancel.

Функция cancel закрывает Done-канал, отменяет дочерние контексты и удаляет связь между родительским и дочерним контекстами. Также при первой отмене она устанавливает причину отмены.

Важно отметить, что при отмене дочернего контекста его родительский контекст не будет отменен, и родительскому контексту никак не сообщается об отмене дочерних контекстов.

Особенности реализации ctx.Value()

Получить значение контекста можно при помощи метода Value. Он принимает ключ и возвращает его значение. Если ключ не существует в текущем контексте, то поиск будет рекурсивно подниматься по дереву контекстов, пока не будет найден корневой контекст. Если значение не найдено, то вернется nil.

Особенность реализации здесь в том, что поиск значения происходит во всех родительских контекстах, — это значительно упрощает его получение. При этом неважно, насколько глубоко вложенным является контекст с искомым значением.

Следует избегать ситуаций с созданием нескольких контекстов с одинаковыми ключами с помощью WithValue(). В таких ситуациях ctx.Value() будет возвращать ближайшее к ctx значение, так как значения в родительских контекстах будут не перезаписаны, а скрыты, и, следовательно, это может стать причиной ошибки.

Кроме того, стоит отметить, что при большой вложенности линейный перебор связного списка контекстов может быть довольно затратным. Поэтому существует рекомендация не использовать контекст как хранилище данных. В том числе есть известная проблема, что, когда цепочка вложенных контекстов достигает определенного размера, это провоцирует довольно дорогостоящую операцию по изменению размера стека горутин, что может отражаться на производительности.

Для менее затратного получения значений из большой цепочки вложенных контекстов может быть использована библиотека zerosnake0/goctx.

Немного практики

Итак, мы разобрали пакет context и теперь предлагаем вам оценить свои знания. 

package main

import "time"

// Вводные:
// Функция executeTask может зависнуть.
// В ней не предусмотрен механизм отмены.
// Она не принимает Context или канал с событием отмены как аргумент.

func main() {

	// Задача:
	// Для функции executeTask написать обертку executeTaskWithTimeout.
	// Функция executeTaskWithTimeout принимает аргументом тайм-аут,
	// через который функция executeTask будет отменена.
	// Если executeTask была отменена по тайм-ауту, нужно вернуть ошибку

	executeTask()
}

func executeTask() {
	time.Sleep(10 * time.Second)
}

РешениеОпределяем функцию:

func executeTaskWithTimeout(ctx context.Context, timeout time.Duration) error {
	// создать контекст с отменой по тайм-ауту
	// реализовать запуск executeTask
	// реализовать возврат ошибки, если executeTask была отменена по таймауту
}

Определяем значение тайм-аута и меняем вызов в main:

func main() {
	timeout := 5 * time.Second
	err := executeTaskWithTimeout(context.Background(), timeout)
	if err != nil {
		log.Fatal(err)
	}
}

Создаем контекст с отменой по тайм-ауту:

func executeTaskWithTimeout(ctx context.Context, timeout time.Duration) error {
	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// реализовать запуск executeTask
	// реализовать возврат ошибки, если executeTask была отменена по тайм-ауту
}

Запускаем функцию в горутине и по выполнении сообщаем в канал done:

func executeTaskWithTimeout(ctx context.Context, timeout time.Duration) error {
	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// создаем буферизированный канал
	done := make(chan struct{}, 1)
	
	// если канал не буферизированный и функция отменена по тайм-ауту,
	// то done <- struct{}{} в горутине заблокируется навсегда и горутина останется зависшей

	go func() {
		executeTask()

		// сигнализируем об окончании executeTask
		done <- struct{}{}

		// закрываем канал
		close(done)

		// закрывать канал важно именно в этой горутине, если закрывать в другом месте,
		// то done <- struct{}{} сгенерирует панику на записи в закрытый канал
	}()
}

Ожидаем выполнение executeTask или отмены контекста по таймауту

func executeTaskWithTimeout(ctx context.Context, timeout time.Duration) error {
	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	done := make(chan struct{}, 1)

	go func() {
		executeTask()
		done <- struct{}{}
		close(done)
	}()

	// ожидаем выполнение executeTask или отмены контекста по тайм-ауту
	select {
	case <-done:  // executeTask выполнилась
		return nil
	case <-timeoutCtx.Done():  // произошла отмена контекста по тайм-ауту
		return timeoutCtx.Err()  // возвращаем ошибку
	}
}

И вот как выглядит итоговое решение:

package main

import (
	"context"
	"log"
	"time"
)

func main() {
	timeout := 5 * time.Second
	err := executeTaskWithTimeout(context.Background(), timeout)
	if err != nil {
		log.Fatal(err)
	}
}

func executeTaskWithTimeout(ctx context.Context, timeout time.Duration) error {
	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	done := make(chan struct{}, 1)

	go func() {
		executeTask()
		done <- struct{}{}
		close(done)
	}()

	select {
	case <-done:
		return nil
	case <-timeoutCtx.Done():
		return timeoutCtx.Err()
	}
}

func executeTask() {
	time.Sleep(10 * time.Second)
}

Итог

Допустим, ты хочешь научиться мастерски владеть инструментом. Стратегий достижения мастерства множество, но одна из самых простых — это учиться у действующих мастеров. Отлично, если это по совместительству мастера, которые создали инструмент. Еще лучше, если это мастера, которые создали его под себя. Да, скорее всего, в последнем случае инструмент не будет выглядеть модно, но зато он будет качественно решать поставленные перед ним задачи.

Go — инструмент, который для себя создали профессионалы. Поэтому реализацию языка и стандартной библиотеки можно и нужно читать. Не потому, что можно использовать какие-нибудь особенности реализации, а потому, что она хорошо написана и там действительно есть чему поучиться.