Case Study: Notification System at Scale

package main

import (
	"container/heap"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"sync"
	"syscall"
	"time"
)

// ===========================================
// 1. TYPES
// ===========================================
type Channel string
type Priority int
type Status string

const (
	Push  Channel = "push"
	Email Channel = "email"
	SMS   Channel = "sms"
	InApp Channel = "in_app"

	Critical Priority = 0
	High     Priority = 1
	Normal   Priority = 2
	Low      Priority = 3
)

func parsePriority(s string) Priority {
	switch s {
	case "critical": return Critical
	case "high": return High
	case "low": return Low
	default: return Normal
	}
}

type Notification struct {
	ID         string            `json:"id"`
	UserID     string            `json:"userId"`
	Type       string            `json:"type"`
	Title      string            `json:"title"`
	Body       string            `json:"body"`
	Channel    Channel           `json:"channel"`
	Priority   Priority          `json:"priority"`
	Status     string            `json:"status"`
	Metadata   map[string]string `json:"metadata"`
	CreatedAt  time.Time         `json:"createdAt"`
	SentAt     *time.Time        `json:"sentAt,omitempty"`
	RetryCount int               `json:"retryCount"`
	index      int
}

type UserPreferences struct {
	UserID          string               `json:"userId"`
	Channels        map[string][]Channel `json:"channels"`
	QuietHoursStart *int                 `json:"quietHoursStart"`
	QuietHoursEnd   *int                 `json:"quietHoursEnd"`
	MaxPerHour      int                  `json:"maxPerHour"`
}

// ===========================================
// 2. PRIORITY QUEUE
// ===========================================
type NotifHeap []*Notification

func (h NotifHeap) Len() int           { return len(h) }
func (h NotifHeap) Less(i, j int) bool { return h[i].Priority < h[j].Priority }
func (h NotifHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}
func (h *NotifHeap) Push(x interface{}) {
	n := x.(*Notification)
	n.index = len(*h)
	*h = append(*h, n)
}
func (h *NotifHeap) Pop() interface{} {
	old := *h
	n := len(old)
	item := old[n-1]
	*h = old[:n-1]
	return item
}

// ===========================================
// 3. RATE LIMITER
// ===========================================
type RateLimiter struct {
	mu      sync.Mutex
	windows map[string][]int64
}

func NewRateLimiter() *RateLimiter {
	return &RateLimiter{windows: make(map[string][]int64)}
}

func (rl *RateLimiter) IsAllowed(key string, maxReq int, windowMs int64) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	now := time.Now().UnixMilli()
	var valid []int64
	for _, t := range rl.windows[key] {
		if t > now-windowMs { valid = append(valid, t) }
	}
	valid = append(valid, now)
	rl.windows[key] = valid
	return len(valid) <= maxReq
}

// ===========================================
// 4. CHANNEL HANDLERS
// ===========================================
type ChannelHandler interface {
	Send(n *Notification) bool
}

type PushHandler struct{}
func (h *PushHandler) Send(n *Notification) bool {
	log.Printf("[PUSH] → %s: %s", n.UserID, n.Title)
	return true
}

type EmailHandler struct{}
func (h *EmailHandler) Send(n *Notification) bool {
	log.Printf("[EMAIL] → %s: %s - %s", n.UserID, n.Title, n.Body)
	return true
}

type SmsHandler struct{}
func (h *SmsHandler) Send(n *Notification) bool {
	log.Printf("[SMS] → %s: %s", n.UserID, n.Body)
	return true
}

type InAppHandler struct {
	mu    sync.Mutex
	store map[string][]*Notification
}

func NewInAppHandler() *InAppHandler {
	return &InAppHandler{store: make(map[string][]*Notification)}
}

func (h *InAppHandler) Send(n *Notification) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.store[n.UserID] = append(h.store[n.UserID], n)
	log.Printf("[IN-APP] → %s: %s", n.UserID, n.Title)
	return true
}

func (h *InAppHandler) Get(userID string) []*Notification {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.store[userID]
}

// ===========================================
// 5. PREFERENCES MANAGER
// ===========================================
type PreferencesManager struct {
	mu    sync.RWMutex
	prefs map[string]*UserPreferences
}

func NewPreferencesManager() *PreferencesManager {
	return &PreferencesManager{prefs: make(map[string]*UserPreferences)}
}

func (pm *PreferencesManager) Get(userID string) *UserPreferences {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	if p, ok := pm.prefs[userID]; ok { return p }
	return &UserPreferences{
		UserID:     userID,
		Channels:   map[string][]Channel{"default": {Push, Email, InApp}},
		MaxPerHour: 50,
	}
}

func (pm *PreferencesManager) Set(userID string, p *UserPreferences) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	pm.prefs[userID] = p
}

func (pm *PreferencesManager) GetChannels(userID, notifType string) []Channel {
	p := pm.Get(userID)
	if ch, ok := p.Channels[notifType]; ok { return ch }
	if ch, ok := p.Channels["default"]; ok { return ch }
	return []Channel{InApp}
}

func (pm *PreferencesManager) IsQuietHours(userID string) bool {
	p := pm.Get(userID)
	if p.QuietHoursStart == nil || p.QuietHoursEnd == nil { return false }
	hour := time.Now().Hour()
	if *p.QuietHoursStart < *p.QuietHoursEnd {
		return hour >= *p.QuietHoursStart && hour < *p.QuietHoursEnd
	}
	return hour >= *p.QuietHoursStart || hour < *p.QuietHoursEnd
}

// ===========================================
// 6. NOTIFICATION SERVICE
// ===========================================
type NotificationService struct {
	mu       sync.Mutex
	queue    NotifHeap
	rl       *RateLimiter
	prefs    *PreferencesManager
	inApp    *InAppHandler
	handlers map[Channel]ChannelHandler
	logs     sync.Map
	stopCh   chan struct{}
}

func NewNotificationService() *NotificationService {
	inApp := NewInAppHandler()
	svc := &NotificationService{
		rl:    NewRateLimiter(),
		prefs: NewPreferencesManager(),
		inApp: inApp,
		handlers: map[Channel]ChannelHandler{
			Push: &PushHandler{}, Email: &EmailHandler{},
			SMS: &SmsHandler{}, InApp: inApp,
		},
		stopCh: make(chan struct{}),
	}
	heap.Init(&svc.queue)
	go svc.processLoop()
	return svc
}

func (svc *NotificationService) Send(userID, notifType, title, body string, priority Priority, metadata map[string]string) []string {
	channels := svc.prefs.GetChannels(userID, notifType)
	var ids []string
	for _, ch := range channels {
		if priority != Critical && svc.prefs.IsQuietHours(userID) { continue }
		n := &Notification{
			ID: fmt.Sprintf("%d", time.Now().UnixNano()),
			UserID: userID, Type: notifType, Title: title, Body: body,
			Channel: ch, Priority: priority, Status: "pending",
			Metadata: metadata, CreatedAt: time.Now().UTC(),
		}
		svc.logs.Store(n.ID, n)
		svc.mu.Lock()
		heap.Push(&svc.queue, n)
		svc.mu.Unlock()
		ids = append(ids, n.ID)
	}
	return ids
}

func (svc *NotificationService) FanOut(userIDs []string, notifType, title, body string, priority Priority) int {
	count := 0
	for _, uid := range userIDs {
		ids := svc.Send(uid, notifType, title, body, priority, nil)
		count += len(ids)
	}
	return count
}

func (svc *NotificationService) processLoop() {
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C: svc.processBatch()
		case <-svc.stopCh: return
		}
	}
}

func (svc *NotificationService) processBatch() {
	svc.mu.Lock()
	batch := svc.queue.Len()
	if batch > 50 { batch = 50 }
	items := make([]*Notification, 0, batch)
	for i := 0; i < batch; i++ {
		items = append(items, heap.Pop(&svc.queue).(*Notification))
	}
	svc.mu.Unlock()

	for _, n := range items {
		key := fmt.Sprintf("%s:%s", n.UserID, n.Channel)
		prefs := svc.prefs.Get(n.UserID)
		if !svc.rl.IsAllowed(key, prefs.MaxPerHour, 3600000) {
			n.Status = "failed"
			log.Printf("[RATE-LIMITED] %s on %s", n.UserID, n.Channel)
			continue
		}
		if svc.handlers[n.Channel].Send(n) {
			n.Status = "delivered"
			now := time.Now().UTC()
			n.SentAt = &now
		} else {
			n.RetryCount++
			if n.RetryCount < 3 {
				delay := time.Duration(math.Pow(2, float64(n.RetryCount))) * time.Second
				go func(notif *Notification) {
					time.Sleep(delay)
					svc.mu.Lock()
					heap.Push(&svc.queue, notif)
					svc.mu.Unlock()
				}(n)
			} else {
				n.Status = "failed"
			}
		}
	}
}

// ===========================================
// 7. HTTP SERVER
// ===========================================
func writeJSON(w http.ResponseWriter, status int, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(data)
}

func main() {
	svc := NewNotificationService()
	notifPattern := regexp.MustCompile(`^/api/notifications/([^/]+)$`)
	prefPattern := regexp.MustCompile(`^/api/preferences/([^/]+)$`)

	mux := http.NewServeMux()

	mux.HandleFunc("/api/notify", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, 405, map[string]string{"error": "Method not allowed"}); return
		}
		var body struct {
			UserID   string            `json:"userId"`
			Type     string            `json:"type"`
			Title    string            `json:"title"`
			Body     string            `json:"body"`
			Priority string            `json:"priority"`
			Metadata map[string]string `json:"metadata"`
		}
		if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil {
			writeJSON(w, 400, map[string]string{"error": "Invalid JSON"}); return
		}
		if body.UserID == "" || body.Title == "" || body.Body == "" {
			writeJSON(w, 400, map[string]string{"error": "userId, title, and body required"}); return
		}
		t := body.Type; if t == "" { t = "default" }
		ids := svc.Send(body.UserID, t, body.Title, body.Body, parsePriority(body.Priority), body.Metadata)
		writeJSON(w, 201, map[string]interface{}{"notificationIds": ids})
	})

	mux.HandleFunc("/api/notify/batch", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeJSON(w, 405, map[string]string{"error": "Method not allowed"}); return
		}
		var body struct {
			UserIDs  []string `json:"userIds"`
			Type     string   `json:"type"`
			Title    string   `json:"title"`
			Body     string   `json:"body"`
			Priority string   `json:"priority"`
		}
		if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil {
			writeJSON(w, 400, map[string]string{"error": "Invalid JSON"}); return
		}
		if len(body.UserIDs) == 0 || body.Title == "" || body.Body == "" {
			writeJSON(w, 400, map[string]string{"error": "userIds, title, and body required"}); return
		}
		count := svc.FanOut(body.UserIDs, body.Type, body.Title, body.Body, parsePriority(body.Priority))
		writeJSON(w, 201, map[string]interface{}{"totalQueued": count})
	})

	mux.HandleFunc("/api/notifications/", func(w http.ResponseWriter, r *http.Request) {
		m := notifPattern.FindStringSubmatch(r.URL.Path)
		if m == nil { writeJSON(w, 404, map[string]string{"error": "Not found"}); return }
		notifs := svc.inApp.Get(m[1])
		writeJSON(w, 200, map[string]interface{}{"notifications": notifs, "count": len(notifs)})
	})

	mux.HandleFunc("/api/preferences/", func(w http.ResponseWriter, r *http.Request) {
		m := prefPattern.FindStringSubmatch(r.URL.Path)
		if m == nil { writeJSON(w, 404, map[string]string{"error": "Not found"}); return }
		if r.Method == http.MethodGet {
			writeJSON(w, 200, svc.prefs.Get(m[1])); return
		}
		if r.Method == http.MethodPut {
			var body UserPreferences
			json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body)
			body.UserID = m[1]
			svc.prefs.Set(m[1], &body)
			writeJSON(w, 200, svc.prefs.Get(m[1])); return
		}
		writeJSON(w, 405, map[string]string{"error": "Method not allowed"})
	})

	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		writeJSON(w, 200, map[string]string{"status": "ok"})
	})

	port := os.Getenv("PORT"); if port == "" { port = "3000" }
	srv := &http.Server{Addr: ":" + port, Handler: mux, ReadTimeout: 5 * time.Second, WriteTimeout: 10 * time.Second}

	go func() {
		log.Printf("Notification Service on http://localhost:%s", port)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed { log.Fatal(err) }
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	log.Println("Shutting down...")
	svc.stopCh <- struct{}{}
	srv.Close()
}