package main

import (
	"context"
	"io"
	"time"

	"golang.org/x/time/rate"
)

type throttledWriter struct {
	ctx     context.Context
	w       io.Writer
	limiter *rate.Limiter
}

func (tw throttledWriter) Write(b []byte) (int, error) {
	if tw.limiter != nil {
		start := time.Now()
		if block := tw.limiter.Burst(); len(b) > block {
			b = b[:block]
		}
		if err := tw.limiter.WaitN(tw.ctx, len(b)); err != nil {
			return 0, err
		}
	}
	return tw.w.Write(b)
}