GoでGraphQL Subscriptionsを実装する

はじめに

この記事はGo3 Advent Calendar 2020の11日目の記事です。gqlgenを使ってGraphQL Subscriptionsを実装する方法とハマったポイントを紹介したいと思います。

利用技術

  • gqlgen
    • GraphQL SchemaからGoのコードを出力するコードファーストなライブラリ
  • Redis Pub/Sub
    • Redisでクライアント間通信を行う
  • echo
    • Webフレームワーク

今回はドメイン3層とgqlgenの生成物を合わせた構成にしています。 詳しくは先日の記事にまとめています。

実装した内容はGitHubに公開しました。
※記事の中のコードはわかりやすさを重視するため今回の内容に関わらない部分は省略しています

github.com

Docker環境の構築

docker-compose.yaml にRedisコンテナを追加します。

version: '3.7'
services:
  redis:
    image: redis:latest
    ports:
      - "6379:6379"

Redis Clientの選定

Goのパッケージに公開されているRedis Clientは公式ドキュメントにまとめられていて、星マークがついてるライブラリがRedis公式に推奨されてます。最初はその中のRadixを使用していましたが、Redisのコマンドを文字列で入力する必要があったため、Redisコマンドをメソッドとして抽象化しているgo-redis/redisを使用しました。

Redis Clientセットアップ

まずライブラリをインストールします。

$ go get -u github.com/go-redis/redis/v8

Redis Clientはインフラ層に定義し、環境変数からURLを取得して初期化しています。

type Store struct {
    Client redis.UniversalClient
    TTL    time.Duration
}

func NewRedisClient(ctx context.Context) (*Store, error) {
    client := redis.NewClient(&redis.Options{
        Addr: os.Getenv("REDIS_URL"),
    })
    err := client.Ping(ctx).Err()
    if err != nil {
        return nil, err
    }
    defaultTTL := 24 * time.Hour
    return &Store{Client: client, TTL: defaultTTL}, nil
}

Pub/SubのクライアントはデータストアであるRedis本体とは目的が異なるため、Storeとは別の構造体で定義しました。

model.Matchドメイン層の構造体で、ユーザー間のマッチングをやり取りするためのデータです。model.Match が生成されたタイミングで通知するchannelを定義します。今回は一つのみにしていますが、GraphQL Subscriptionsが複数ある場合はchannelも複数定義していきます。

type PubSubClient struct {
    Store         Store
    MatchChannels map[string]chan *model.Match
    Mutex         sync.Mutex
}

func NewPubSubClient(ctx context.Context, store Store) *PubSubClient {
    subscriber := &PubSubClient{
        Store:         store,
        MatchChannels: map[string]chan *model.Match{},
        Mutex:         sync.Mutex{},
    }
    return subscriber
}

定義した2つのクライアントをResolverに注入します。

type Resolver struct {
    DB           *gorm.DB
    StorePool    *store.Store
    PubSubClient *store.PubSubClient
}

各ResolverでRedisを利用できるようになりました。

スキーマの追加

GraphQLスキーマにSubscriptionsを追加します。 今回はユーザーとユーザー間を Match で関連付けるデータ構造にしているため、 Match オブジェクトが生成されたタイミングでユーザーに通知する仕組みを作ります。

type Mutation {
    createMatch(input: NewMatch!): Match!
}

type Subscription {
    createMatch(userUID: String!): Match
}

スキーマを編集した後、 gqlgen コマンドを実行してSubscriptionsのResolverを生成します。 resolver パッケージに以下のコードが追加されたら完了です。

func (r *subscriptionResolver) CreateMatch(ctx context.Context, userUID string) (<-chan *model.Match, error) {
   panic(fmt.Errorf("not implemented"))
}

他のResolverと異なり、関数の返り値がchannelになっています。

Subscriptionsの実装

GraphQLでSubscriptionsを実装する場合、以下の内容を追加する必要があります。

  1. APIがRedis PubSubから配信されるメッセージを受信
  2. ユーザーがSubscriptionsを実行したときびWebSocket通信
  3. ユーザーがMutationを実行したときにRedis PubSubに配信

実装は PubSub Client のメソッドに追加していき、Resolverでは定義されたメソッドを呼び出すようにしました。

それぞれ実装をみていきます。

1. APIがRedis PubSubから配信されるメッセージを受信

API起動時にRedis PubSubからのメッセージを受信し、パースしたデータをChannelに流し込みます。 msg.Payload がRedisからのメッセージになっています。

func (s *PubSubClient) StartMatchChannel(ctx context.Context) {
    go func() {
        pubsub := s.Store.Client.Subscribe(ctx, channel.MatchChannel)
        defer pubsub.Close()

        for {
            msgInterface, _ := pubsub.Receive(ctx)
            switch msg := msgInterface.(type) {
            case *redis.Message:
                m := model.Match{}
                if err := json.Unmarshal([]byte(msg.Payload), &m); err != nil {
                    continue
                }
                for _, ch := range s.MatchChannels {
                    ch <- &m
                }
            default:
            }
        }
    }()
}

2. ユーザーがSubscriptionsを実行したときのWebSocket通信

ユーザーがSubscriptionsを実行するときにユーザーごとのchannelを生成します。

構造体に定義しているchannelのキーをユーザーIDにしてRedisに保存し、受信状態を作ります。

func (s *PubSubClient) Receive(ctx context.Context, userUID string) <-chan *model.Match {
    _, err := s.Store.Client.SetXX(ctx, userUID, userUID, 60*time.Minute).Result()
    if err != nil {
        return nil
    }

    match := make(chan *model.Match, 1)
    s.MatchChannels[userUID] = match

    go func() {
        <-ctx.Done()
        delete(s.MatchChannels, userUID)
        s.Store.Client.Del(ctx, userUID)
    }()

    return match
}

3. ユーザーがMutationを実行したときにRedis PubSubに配信

最後にSubscriptionsが発火するMutation内でRedis PubSubにデータを配信します。

func (s *PubSubClient) Publish(ctx context.Context, match *model.Match) error {
    matchJSON, err := json.Marshal(match)
    if err != nil {
        return err
    }
    s.Store.Client.Publish(ctx, channel.MatchChannel, matchJSON)
    return nil
}

これでMutationが実行されるたびにユーザーにデータを配信できます。

ハマったポイント

gqlgen generateできない

最初にリポジトリを生成してから2ヶ月ほど経過してから修正したため、gqlgenのバージョンがCLIは0.12.2に対してgo.modは0.13.0と差分が出てました。 この状態で gqlgen generate を実行すると generated.go のなかで not declared by package エラーが発生してしまいます。

今回はCLIのgqlgenを再インストールすることで解決しました。

WebSocket通信できない

gqlgenの handler には NewDefaultServer メソッドが存在していて、返り値のserverには大枠の設定が自動で行われています。しかしWebSocketの設定も追加されているためRedis PubSubと接続するときにoriginチェックの処理で落ちてしまいます。

そのため handlerNew メソッドを使って必要な部分だけ Transport するようにしました。

WebSocket通信が認証エラーになる

実装をまとめてるリポジトリでは元々Firebase Authを組み込んでいいて、Go上ではクライアントから渡されたトークンをFirebase Authに確認しユーザーIDを取得する処理を行っていました。 Firebase Authの認証処理をechoの middleware で定義していたため、WebSocket通信でも認証チェックが行われてしまい、ヘッダーのトークンを参照できないためエラーになってしまいました。

今回は echo.Context に含まれている IsWebSocket メソッドを利用してWebSocker通信では認証処理をしないようにしました。

まとめ

gqlgenとRedis PubSubを使ってGraphQL Subscriptionsを実装しました。 Goのchannelに流し込むとクライアントに配信してくれるので比較的簡単に機能を作れました。 今後はchannelが増えたときに設計をどう拡張・修正していくか考えてみたいです。

参考