巧妙利用装饰器模式给WebSocket连接新增持久化

引言

我们在写代码时，常常希望可以做到不修改现有的业务逻辑，优雅地为系统添加新的功能特性！

本文将通过分析一个WebSocket消息持久化的实际案例，深入探讨装饰器模式如何通过接口组合实现功能的"无感"扩展。

设计模式概述

装饰器模式的核心思想

装饰器模式（Decorator Pattern）是一种结构型设计模式，它允许你通过将对象包装在装饰器类的对象中来为原对象添加新的行为，而无需修改原对象的代码。

这种模式遵循了开闭原则（Open/Closed Principle）：对扩展开放，对修改关闭。

接口组合的威力

在Go语言中，接口组合（Interface Composition）是实现装饰器模式的优雅方式。通过将多个小接口组合成一个大接口，我们可以实现功能的模块化和可插拔性。

案例分析：WebSocket消息持久化

让我们来分析一个实际的代码案例，看看如何通过装饰器模式实现WebSocket消息的"无感"持久化。

接口设计

首先，我们定义了一个组合接口：

go 复制代码

type WebsocketBotWithPersistence interface {
    service.Bot
    websockettil.ConnectionReadWriteHandler
    WithPersistence()
}

这个接口巧妙地组合了三个功能：

service.Bot：定义Bot的基本属性（ID和名称）
websockettil.ConnectionReadWriteHandler：定义WebSocket连接的读写处理能力
WithPersistence()：标记该Bot支持持久化功能

装饰器实现

装饰器的核心实现如下：

go 复制代码

type PersistenceDecorator struct {
    bot       WebsocketBotWithPersistence
    persister sdks.MessagePersister
}

func NewPersistenceDecorator(bot WebsocketBotWithPersistence, persister sdks.MessagePersister) *PersistenceDecorator {
    return &PersistenceDecorator{
        bot:       bot,
        persister: persister,
    }
}

func (d *PersistenceDecorator) OnRead(ctx context.Context, conn websockettil.ConnectionInterface, message websockettil.ConnectionMessage) error {
    // 先处理持久化逻辑
    if err := d.persister.ProcessAIMessage(ctx, conn, message.MessageData); err != nil {
        logger.FromContext(ctx).Errorf("failed to process ai message: %v", err)
    }
    
    // 再调用原有逻辑
    return d.bot.OnRead(ctx, conn, message)
}

func (d *PersistenceDecorator) OnWrite(ctx context.Context, conn websockettil.ConnectionInterface, message websockettil.ConnectionMessage) error {
    // 先处理持久化逻辑
    var param sdks.AIServiceBotChatParam
    if err := jsoniter.Unmarshal(message.MessageData, &param); err == nil {
        userInfo, ok := ctxtil.GetYunYingBackendUserInfo(ctx)
        if ok {
            if e := d.persister.PersistUserMessage(ctx, conn, param, userInfo.NamePinYin); e != nil {
                logger.FromContext(ctx).Errorf("failed to persist user message: %v", e)
            }
        }
    }
    
    // 再调用原有逻辑
    return d.bot.OnWrite(ctx, conn, message)
}

"无感"的实现机制

1. 接口透明性

装饰器实现了与原对象相同的接口，这意味着调用者无需知道对象是否被装饰过。无论是原始的Bot对象还是被装饰过的Bot对象，都可以通过相同的接口进行操作。

2. 透明的功能注入

我们可以提前使用 wire 生成初始化的代码，然后在系统初始化时，通过反射机制，自动为支持持久化的Bot添加持久化装饰器！

go 复制代码

func (m *Map) Init() {
    m.BotMap = make(map[string]service.Bot)
    rv := reflect.ValueOf(*m)
    for i := 0; i < rv.NumField(); i++ {
        a, ok := rv.Field(i).Interface().(service.Bot)
        if !ok {
            continue
        }
        // 关键：自动判断是否需要持久化装饰器
        if decorator, ok2 := a.(bots.WebsocketBotWithPersistence); ok2 {
            // 套娃
            m.BotMap[a.BotID()] = bots.NewPersistenceDecorator(decorator, m.MessagePersister)
            continue
        }
        // 不支持
        m.BotMap[a.BotID()] = a
    }
}

3. 链式调用保证

装饰器在执行自己的逻辑后，总是调用被装饰对象的对应方法，保证了原有业务逻辑的完整性。

持久化机制的技术细节

消息流转过程

用户消息处理 ：在OnWrite方法中，装饰器解析用户消息并调用PersistUserMessage进行持久化
AI消息处理 ：在OnRead方法中，装饰器调用ProcessAIMessage处理AI回复消息
缓冲区机制：使用连接的metadata存储消息缓冲区，实现流式消息的组装

go 复制代码

// 确保对话存在，如果不存在则创建
func (mp *MessagePersister) ensureConversationExists(ctx context.Context, conn websockettil.ConnectionInterface, param sdks.AIServiceBotChatParam, creator string) error {
    // 检查连接metadata中是否已经标记对话存在
    conversationKey := "conversation_ensured_" + param.HistoryID
    if _, exist := conn.GetMetadata(conversationKey); exist {
        return nil
    }
    
    // 持久化对话的逻辑...
    
    // 标记对话已存在
    conn.SetMetadata(conversationKey, true)
    return nil
}

使用装饰器的好处

1. 单一职责原则

原始Bot：专注于业务逻辑处理
装饰器：专注于持久化功能
消息持久化器：专注于数据存储逻辑

go 复制代码

if decorator, ok := bot.(bots.WebsocketBotWithPersistence); ok {
    // 该Bot支持持久化，自动添加装饰器
    bot = bots.NewPersistenceDecorator(decorator, persister)
}

2. 开闭原则

对扩展开放：可以轻松添加新的装饰器（如日志记录、性能监控等），做到多层装饰（套娃大法）
对修改关闭：无需修改原有Bot的代码

go 复制代码

// 套娃+1
bot = NewPersistenceDecorator(bot, persister)
// 套娃+2
bot = NewLoggingDecorator(bot, logger)
// 套娃+3
bot = NewMetricsDecorator(bot, metrics)
// 套娃+10086
//
// 我还是从前那个bot(少年)没有一丝丝改变......

3. 依赖注入友好

装饰器模式与依赖注入框架完美结合，通过自动装配实现功能的自动组装。

4. 测试友好

每个组件都可以独立测试：

Bot可以独立测试业务逻辑
装饰器可以通过Mock对象测试
持久化器可以独立测试存储逻辑

5. 渐进式迁移

对于现有系统，可以渐进式地添加新功能，无需大规模重构。

总结

装饰器模式通过接口组合实现的"无感"持久化具有以下特点：

透明性：调用者无需知道对象是否被装饰
可扩展性：可以轻松添加新的功能装饰器
可维护性：各个组件职责清晰，易于维护
可测试性：各个组件可以独立测试

这种设计方式不仅体现了面向对象设计的核心原则，也展示了Go语言接口组合的强大威力。通过合理使用装饰器模式，我们可以构建出既灵活又可维护的系统架构。