|
| 1 | +(ns cmr.message-queue.queue.sqs-v2 |
| 2 | + "AWS SDK v2 implementation of the CMR Queue protocol." |
| 3 | + (:require |
| 4 | + [cheshire.core :as json] |
| 5 | + [clojure.core.async :as async] |
| 6 | + [clojure.string :as string] |
| 7 | + [cmr.common.config :refer [defconfig]] |
| 8 | + [cmr.common.dev.record-pretty-printer :as record-pretty-printer] |
| 9 | + [cmr.common.lifecycle :as lifecycle] |
| 10 | + [cmr.common.log :refer [debug error info warn]] |
| 11 | + [cmr.message-queue.queue.names :as names] |
| 12 | + [cmr.message-queue.queue.queue-protocol :as queue-protocol]) |
| 13 | + (:import |
| 14 | + (java.net URI) |
| 15 | + (software.amazon.awssdk.auth.credentials AwsBasicCredentials StaticCredentialsProvider) |
| 16 | + (software.amazon.awssdk.regions Region) |
| 17 | + (software.amazon.awssdk.services.sns SnsClient) |
| 18 | + (software.amazon.awssdk.services.sns.model CreateTopicRequest ListSubscriptionsByTopicRequest |
| 19 | + ListTopicsRequest PublishRequest |
| 20 | + SetSubscriptionAttributesRequest SubscribeRequest) |
| 21 | + (software.amazon.awssdk.services.sqs SqsClient) |
| 22 | + (software.amazon.awssdk.services.sqs.model CreateQueueRequest DeleteMessageRequest |
| 23 | + GetQueueAttributesRequest GetQueueUrlRequest |
| 24 | + PurgeQueueRequest QueueAttributeName |
| 25 | + ReceiveMessageRequest SendMessageRequest |
| 26 | + SetQueueAttributesRequest))) |
| 27 | + |
| 28 | +(defconfig queue-polling-timeout |
| 29 | + "Seconds to wait for an SQS long poll." |
| 30 | + {:default 20 :type Long}) |
| 31 | +(defconfig default-queue-visibility-timeout |
| 32 | + "Default visibility timeout." |
| 33 | + {:default 300 :type Long}) |
| 34 | +(defconfig provider-queue-visibility-timeout |
| 35 | + "Provider queue visibility timeout." |
| 36 | + {:default 43200 :type Long}) |
| 37 | +(defconfig default-num-tries |
| 38 | + "Default receive attempts before redrive." |
| 39 | + {:default 5 :type Long}) |
| 40 | +(defconfig sqs-endpoint |
| 41 | + "Optional SQS endpoint override." |
| 42 | + {:default nil :type String}) |
| 43 | +(defconfig sns-endpoint |
| 44 | + "Optional SNS endpoint override." |
| 45 | + {:default nil :type String}) |
| 46 | +(defconfig sqs-extend-policy-remaining-exchanges |
| 47 | + "Whether additional topic bindings extend the queue policy." |
| 48 | + {:default true :type Boolean}) |
| 49 | + |
| 50 | +(defn queue-visibility-timeout [queue-name] |
| 51 | + (if (string/includes? queue-name "provider") |
| 52 | + (provider-queue-visibility-timeout) |
| 53 | + (default-queue-visibility-timeout))) |
| 54 | + |
| 55 | +(defn dead-letter-queue [normalized-name] |
| 56 | + (str normalized-name "_dead_letter_queue")) |
| 57 | + |
| 58 | +(defn arn->name [arn] |
| 59 | + (string/replace arn #".*:" "")) |
| 60 | + |
| 61 | +(defn subscription-endpoint->name [endpoint] |
| 62 | + (if (string/starts-with? endpoint "http") |
| 63 | + (last (string/split endpoint #"/")) |
| 64 | + (arn->name endpoint))) |
| 65 | + |
| 66 | +(defn configure-builder |
| 67 | + "Applies an endpoint override and the fixed local region when endpoint is configured." |
| 68 | + [builder endpoint] |
| 69 | + (if endpoint |
| 70 | + (-> builder |
| 71 | + (.endpointOverride (URI/create endpoint)) |
| 72 | + (.region Region/US_EAST_1) |
| 73 | + (.credentialsProvider |
| 74 | + (StaticCredentialsProvider/create (AwsBasicCredentials/create "local" "local")))) |
| 75 | + builder)) |
| 76 | + |
| 77 | +(defn create-aws-client [type] |
| 78 | + (case type |
| 79 | + :sqs (.build (configure-builder (SqsClient/builder) (sqs-endpoint))) |
| 80 | + :sns (.build (configure-builder (SnsClient/builder) (sns-endpoint))))) |
| 81 | + |
| 82 | +(defn queue-url [^SqsClient client normalized-name] |
| 83 | + (.queueUrl (.getQueueUrl client (-> (GetQueueUrlRequest/builder) |
| 84 | + (.queueName normalized-name) .build)))) |
| 85 | + |
| 86 | +(defn queue-arn [^SqsClient client url] |
| 87 | + (let [response (.getQueueAttributes |
| 88 | + client (-> (GetQueueAttributesRequest/builder) |
| 89 | + (.queueUrl url) |
| 90 | + (.attributeNames [QueueAttributeName/QUEUE_ARN]) |
| 91 | + .build))] |
| 92 | + (get (.attributes response) QueueAttributeName/QUEUE_ARN))) |
| 93 | + |
| 94 | +(defn create-queue! |
| 95 | + "Idempotently creates a queue and DLQ, then applies visibility and redrive attributes." |
| 96 | + [^SqsClient client queue-name max-tries visibility-timeout] |
| 97 | + (let [name (names/normalize-queue-name queue-name) |
| 98 | + dlq-name (dead-letter-queue name) |
| 99 | + dlq-url (.queueUrl (.createQueue client (-> (CreateQueueRequest/builder) |
| 100 | + (.queueName dlq-name) .build))) |
| 101 | + dlq-arn (queue-arn client dlq-url) |
| 102 | + url (.queueUrl (.createQueue client (-> (CreateQueueRequest/builder) |
| 103 | + (.queueName name) .build))) |
| 104 | + attrs {QueueAttributeName/VISIBILITY_TIMEOUT (str visibility-timeout) |
| 105 | + QueueAttributeName/REDRIVE_POLICY |
| 106 | + (json/generate-string {:maxReceiveCount (str max-tries) |
| 107 | + :deadLetterTargetArn dlq-arn})}] |
| 108 | + (.setQueueAttributes client (-> (SetQueueAttributesRequest/builder) |
| 109 | + (.queueUrl url) (.attributes attrs) .build)) |
| 110 | + url)) |
| 111 | + |
| 112 | +(defn create-topic! [^SnsClient client exchange-name] |
| 113 | + (.topicArn (.createTopic client (-> (CreateTopicRequest/builder) |
| 114 | + (.name (names/normalize-queue-name exchange-name)) .build)))) |
| 115 | + |
| 116 | +(defn queue-policy [queue-arn topic-arns] |
| 117 | + (json/generate-string |
| 118 | + {:Version "2012-10-17" |
| 119 | + :Statement (mapv (fn [topic-arn] |
| 120 | + {:Sid (str "Allow-" (Math/abs (hash topic-arn))) |
| 121 | + :Effect "Allow" |
| 122 | + :Principal {:Service "sns.amazonaws.com"} |
| 123 | + :Action "SQS:SendMessage" |
| 124 | + :Resource queue-arn |
| 125 | + :Condition {:ArnEquals {"aws:SourceArn" topic-arn}}}) |
| 126 | + topic-arns)})) |
| 127 | + |
| 128 | +(defn bind-queue-to-exchanges! |
| 129 | + [^SnsClient sns ^SqsClient sqs exchange-names queue-name] |
| 130 | + (let [url (queue-url sqs (names/normalize-queue-name queue-name)) |
| 131 | + q-arn (queue-arn sqs url) |
| 132 | + topic-arns (mapv #(create-topic! sns %) exchange-names) |
| 133 | + policy-topics (if (sqs-extend-policy-remaining-exchanges) |
| 134 | + topic-arns |
| 135 | + (take-last 1 topic-arns))] |
| 136 | + (.setQueueAttributes sqs (-> (SetQueueAttributesRequest/builder) |
| 137 | + (.queueUrl url) |
| 138 | + (.attributes {QueueAttributeName/POLICY |
| 139 | + (queue-policy q-arn policy-topics)}) |
| 140 | + .build)) |
| 141 | + (doseq [topic-arn topic-arns] |
| 142 | + (let [sub-arn (.subscriptionArn |
| 143 | + (.subscribe sns (-> (SubscribeRequest/builder) |
| 144 | + (.topicArn topic-arn) (.protocol "sqs") |
| 145 | + (.endpoint q-arn) .build)))] |
| 146 | + (.setSubscriptionAttributes |
| 147 | + sns (-> (SetSubscriptionAttributesRequest/builder) |
| 148 | + (.subscriptionArn sub-arn) (.attributeName "RawMessageDelivery") |
| 149 | + (.attributeValue "true") .build)))))) |
| 150 | + |
| 151 | +(defn get-topic-arn [^SnsClient client exchange-name] |
| 152 | + (let [wanted (names/normalize-queue-name exchange-name)] |
| 153 | + (loop [token nil] |
| 154 | + (let [builder (ListTopicsRequest/builder) |
| 155 | + _ (when token (.nextToken builder token)) |
| 156 | + response (.listTopics client (.build builder)) |
| 157 | + found (some #(when (= wanted (arn->name (.topicArn %))) (.topicArn %)) |
| 158 | + (.topics response))] |
| 159 | + (or found (when-let [next-token (.nextToken response)] (recur next-token))))))) |
| 160 | + |
| 161 | +(defn create-async-handler |
| 162 | + ([broker queue-name handler] (create-async-handler broker queue-name handler true)) |
| 163 | + ([broker queue-name handler auto-reconnect?] |
| 164 | + (info "Starting listener for queue:" queue-name) |
| 165 | + (let [name (names/normalize-queue-name queue-name) |
| 166 | + url (queue-url @(:sqs-client-atom broker) name) |
| 167 | + request (-> (ReceiveMessageRequest/builder) (.queueUrl url) |
| 168 | + (.maxNumberOfMessages (int 1)) |
| 169 | + (.waitTimeSeconds (int (queue-polling-timeout))) .build)] |
| 170 | + (async/thread |
| 171 | + (loop [] |
| 172 | + (try |
| 173 | + (when-let [message (first (.messages (.receiveMessage |
| 174 | + ^SqsClient @(:sqs-client-atom broker) request)))] |
| 175 | + (try |
| 176 | + (handler (json/decode (.body message) true)) |
| 177 | + (.deleteMessage ^SqsClient @(:sqs-client-atom broker) |
| 178 | + (-> (DeleteMessageRequest/builder) (.queueUrl url) |
| 179 | + (.receiptHandle (.receiptHandle message)) .build)) |
| 180 | + (catch Throwable e |
| 181 | + (error e "Message processing failed for message" (pr-str message) |
| 182 | + "on queue" name)))) |
| 183 | + (catch Throwable e |
| 184 | + (if (= "cmr.message_queue.test.ExitException" (.getName (class e))) |
| 185 | + (do |
| 186 | + (error "Async handler for queue" name "exiting.") |
| 187 | + (throw e)) |
| 188 | + (do |
| 189 | + (error e "Async handler for queue" name "continuing after failed receive.") |
| 190 | + (Thread/sleep 1000) |
| 191 | + (when auto-reconnect? |
| 192 | + (warn "Recreating SQS v2 client.") |
| 193 | + (let [old @(:sqs-client-atom broker) |
| 194 | + replacement (create-aws-client :sqs)] |
| 195 | + (reset! (:sqs-client-atom broker) replacement) |
| 196 | + (.close ^SqsClient old))))))) |
| 197 | + (recur)))))) |
| 198 | + |
| 199 | +(defrecord SQSQueueBrokerV2 |
| 200 | + [sns-client-atom sqs-client-atom queues normalized-queue-names exchanges |
| 201 | + queues-to-policies queues-to-exchanges] |
| 202 | + lifecycle/Lifecycle |
| 203 | + (start [this _] |
| 204 | + (let [sqs (create-aws-client :sqs) |
| 205 | + sns (create-aws-client :sns)] |
| 206 | + (try |
| 207 | + (doseq [queue queues] |
| 208 | + (create-queue! sqs queue |
| 209 | + (get-in queues-to-policies [queue :max-tries] (default-num-tries)) |
| 210 | + (get-in queues-to-policies [queue :visibility-timeout-secs] |
| 211 | + (queue-visibility-timeout queue)))) |
| 212 | + (doseq [exchange exchanges] (create-topic! sns exchange)) |
| 213 | + (doseq [[queue bound-exchanges] queues-to-exchanges] |
| 214 | + (bind-queue-to-exchanges! sns sqs bound-exchanges queue)) |
| 215 | + (assoc this |
| 216 | + :sqs-client-atom (atom sqs) |
| 217 | + :sns-client-atom (atom sns) |
| 218 | + :normalized-queue-names |
| 219 | + (into {} (map (juxt names/normalize-queue-name identity) queues))) |
| 220 | + (catch Throwable e |
| 221 | + (.close sqs) |
| 222 | + (.close sns) |
| 223 | + (throw e))))) |
| 224 | + (stop [this _] |
| 225 | + (when sns-client-atom (.close ^SnsClient @sns-client-atom)) |
| 226 | + (when sqs-client-atom (.close ^SqsClient @sqs-client-atom)) |
| 227 | + this) |
| 228 | + queue-protocol/Queue |
| 229 | + (publish-to-queue [_ queue-name msg] |
| 230 | + (let [client ^SqsClient @sqs-client-atom |
| 231 | + url (queue-url client (names/normalize-queue-name queue-name))] |
| 232 | + (.sendMessage client (-> (SendMessageRequest/builder) (.queueUrl url) |
| 233 | + (.messageBody (json/generate-string msg)) .build)))) |
| 234 | + (get-queues-bound-to-exchange [_ exchange-name] |
| 235 | + (let [client ^SnsClient @sns-client-atom |
| 236 | + topic-arn (get-topic-arn client exchange-name)] |
| 237 | + (loop [token nil result []] |
| 238 | + (let [builder (-> (ListSubscriptionsByTopicRequest/builder) (.topicArn topic-arn)) |
| 239 | + _ (when token (.nextToken builder token)) |
| 240 | + response (.listSubscriptionsByTopic client (.build builder)) |
| 241 | + names-found (map #(get normalized-queue-names |
| 242 | + (subscription-endpoint->name (.endpoint %)) |
| 243 | + (subscription-endpoint->name (.endpoint %))) |
| 244 | + (.subscriptions response)) |
| 245 | + result (into result names-found)] |
| 246 | + (if-let [next-token (.nextToken response)] |
| 247 | + (recur next-token result) |
| 248 | + result))))) |
| 249 | + (publish-to-exchange [_ exchange-name msg] |
| 250 | + (let [client ^SnsClient @sns-client-atom] |
| 251 | + (.publish client (-> (PublishRequest/builder) |
| 252 | + (.topicArn (get-topic-arn client exchange-name)) |
| 253 | + (.message (json/generate-string msg)) .build)))) |
| 254 | + (subscribe [this queue-name handler] |
| 255 | + (create-async-handler this queue-name handler)) |
| 256 | + (reset [_] |
| 257 | + (let [client ^SqsClient @sqs-client-atom] |
| 258 | + (doseq [queue queues |
| 259 | + name [(names/normalize-queue-name queue) |
| 260 | + (dead-letter-queue (names/normalize-queue-name queue))]] |
| 261 | + (.purgeQueue client (-> (PurgeQueueRequest/builder) |
| 262 | + (.queueUrl (queue-url client name)) .build))))) |
| 263 | + (reconnect [this] |
| 264 | + (warn "Recreating SNS v2 client.") |
| 265 | + (let [old @sns-client-atom |
| 266 | + replacement (create-aws-client :sns)] |
| 267 | + (reset! sns-client-atom replacement) |
| 268 | + (.close ^SnsClient old) |
| 269 | + this)) |
| 270 | + (health [this] |
| 271 | + (try |
| 272 | + (queue-protocol/get-queues-bound-to-exchange this (first exchanges)) |
| 273 | + {:ok? true} |
| 274 | + (catch Throwable e {:ok? false :msg (.getMessage e)})))) |
| 275 | + |
| 276 | +(record-pretty-printer/enable-record-pretty-printing SQSQueueBrokerV2) |
| 277 | + |
| 278 | +(defn create-queue-broker [{:keys [queues exchanges queues-to-policies queues-to-exchanges]}] |
| 279 | + (->SQSQueueBrokerV2 nil nil queues nil exchanges queues-to-policies queues-to-exchanges)) |
0 commit comments